diff --git a/shodan/__main__.py b/shodan/__main__.py index 4093b94..6634c64 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,20 @@ import time import json +try: + from importlib.metadata import PackageNotFoundError, entry_points, version as package_version +except ImportError: + PackageNotFoundError = None + entry_points = None + package_version = None + + try: + import pkg_resources + except ImportError: + pkg_resources = None +else: + pkg_resources = None + # 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 +63,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 +88,38 @@ basestring = str +def iter_plugin_entry_points(group): + if entry_points is not None: + try: + return tuple(entry_points(group=group)) + except TypeError: + discovered = entry_points() + if hasattr(discovered, 'select'): + return tuple(discovered.select(group=group)) + return tuple(discovered.get(group, [])) + + if pkg_resources is not None: + return tuple(pkg_resources.iter_entry_points(group)) + + return tuple() + + +def get_installed_version(): + if package_version is not None: + try: + return package_version('shodan') + except PackageNotFoundError: + pass + + if pkg_resources is not None: + return pkg_resources.get_distribution('shodan').version + + return 'unknown' + + # 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 +983,7 @@ def radar(): @main.command() def version(): """Print version of this tool.""" - print(pkg_resources.get_distribution("shodan").version) + print(get_installed_version()) if __name__ == '__main__': diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..84348e7 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,20 @@ +import importlib +import sys +import unittest +from unittest import mock + + +class CliTests(unittest.TestCase): + + def tearDown(self): + sys.modules.pop('shodan.__main__', None) + + def test_cli_imports_without_pkg_resources(self): + with mock.patch.dict(sys.modules, {'pkg_resources': None}): + cli = importlib.import_module('shodan.__main__') + + self.assertTrue(callable(cli.main)) + + +if __name__ == '__main__': + unittest.main()