diff --git a/requirements.txt b/requirements.txt index 2692414..f5c36d4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ colorama requests>=2.2.1 XlsxWriter ipaddress;python_version<='2.7' -tldextract \ No newline at end of file +importlib-metadata;python_version<'3.8' +tldextract diff --git a/shodan/__main__.py b/shodan/__main__.py index 4093b94..a18cf54 100644 --- a/shodan/__main__.py +++ b/shodan/__main__.py @@ -30,7 +30,6 @@ import csv import os import os.path -import pkg_resources import shodan import shodan.helpers as helpers import threading @@ -38,6 +37,11 @@ import time import json +try: + from importlib.metadata import entry_points, version as package_version +except ImportError: + from importlib_metadata import entry_points, version as package_version + # The file converters that are used to go from .json.gz to various other formats from shodan.cli.converter import CsvConverter, KmlConverter, GeoJsonConverter, ExcelConverter, ImagesConverter @@ -50,7 +54,6 @@ # Allow 3rd-parties to develop custom commands from click_plugins import with_plugins -from pkg_resources import iter_entry_points # Large subcommands are stored in separate modules from shodan.cli.alert import alert @@ -76,9 +79,16 @@ basestring = str +def iter_plugin_entry_points(group): + discovered = entry_points() + if isinstance(discovered, dict): + return discovered.get(group, ()) + return discovered.select(group=group) + + # Define the main entry point for all of our commands # and expose a way for 3rd-party plugins to tie into the Shodan CLI. -@with_plugins(iter_entry_points('shodan.cli.plugins')) +@with_plugins(iter_plugin_entry_points('shodan.cli.plugins')) @click.group(context_settings=CONTEXT_SETTINGS) def main(): pass @@ -942,7 +952,7 @@ def radar(): @main.command() def version(): """Print version of this tool.""" - print(pkg_resources.get_distribution("shodan").version) + print(package_version('shodan')) if __name__ == '__main__': diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..d437485 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,61 @@ +import importlib +import sys +import unittest +from types import ModuleType +from unittest import mock + +from click.testing import CliRunner + + +class CliTests(unittest.TestCase): + + def test_plugins_load_from_legacy_entry_point_mapping(self): + # Given the mapping returned by older importlib-metadata releases + cli = importlib.import_module('shodan.__main__') + plugin = object() + + with mock.patch('shodan.__main__.entry_points', return_value={'shodan.cli.plugins': (plugin,)}): + # When the CLI discovers third-party plugins + plugins = tuple(cli.iter_plugin_entry_points('shodan.cli.plugins')) + + # Then it returns the entries for the requested group + self.assertEqual((plugin,), plugins) + + def test_version_reports_installed_distribution(self): + # Given the installed Shodan CLI + cli = importlib.import_module('shodan.__main__') + + # When its version command is invoked + result = CliRunner().invoke(cli.main, ['version']) + + # Then it reports the installed distribution version + self.assertEqual(0, result.exit_code, result.output) + self.assertRegex(result.output.strip(), r'^\d+\.\d+\.\d+$') + + def test_help_renders_without_pkg_resources(self): + # Given an isolated runtime where pkg_resources is unavailable + previous = sys.modules.get('pkg_resources') + sys.modules['pkg_resources'] = ModuleType('pkg_resources') + sys.modules.pop('shodan.__main__', None) + + try: + # When the CLI is imported and asked to render help + try: + cli = importlib.import_module('shodan.__main__') + except ImportError as exc: + self.fail('CLI import should not require pkg_resources: {}'.format(exc)) + result = CliRunner().invoke(cli.main, ['--help']) + finally: + sys.modules.pop('shodan.__main__', None) + if previous is not None: + sys.modules['pkg_resources'] = previous + else: + sys.modules.pop('pkg_resources', None) + + # Then help renders successfully through the real Click command + self.assertEqual(0, result.exit_code, result.output) + self.assertIn('Usage:', result.output) + + +if __name__ == '__main__': + unittest.main()