From fa8f6be558f1fdf406fc626be413e3ebee3dc8f0 Mon Sep 17 00:00:00 2001
From: Kamil Mankowski <mankowski@cert.at>
Date: Thu, 28 Sep 2023 14:19:36 +0200
Subject: [PATCH] Fix usages of importing bot's module

---
 intelmq/bin/intelmqctl.py            |  2 +-
 intelmq/lib/bot_debugger.py          |  2 +-
 intelmq/lib/utils.py                 | 13 +++++++++++++
 intelmq/tests/bin/test_intelmqctl.py | 12 ++++++++++++
 intelmq/tests/lib/test_utils.py      | 12 ++++++++----
 5 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/intelmq/bin/intelmqctl.py b/intelmq/bin/intelmqctl.py
index 380b27fa1..51301b1d8 100644
--- a/intelmq/bin/intelmqctl.py
+++ b/intelmq/bin/intelmqctl.py
@@ -931,7 +931,7 @@ def check(self, no_connections=False, check_executables=True):
             if bot_id != 'global':
                 # importable module
                 try:
-                    bot_module = importlib.import_module(bot_config['module'])
+                    bot_module = importlib.import_module(utils.get_bot_module_name(bot_config['module']))
                 except ImportError as exc:
                     check_logger.error('Incomplete installation: Bot %r not importable: %r.', bot_id, exc)
                     retval = 1
diff --git a/intelmq/lib/bot_debugger.py b/intelmq/lib/bot_debugger.py
index 4853bebf4..3694a8ec0 100644
--- a/intelmq/lib/bot_debugger.py
+++ b/intelmq/lib/bot_debugger.py
@@ -50,7 +50,7 @@ def __init__(self, runtime_configuration, bot_id, run_subcommand=None, console_t
         self.dryrun = dryrun
         self.msg = msg
         self.show = show
-        module = import_module(self.runtime_configuration['module'])
+        module = import_module(utils.get_bot_module_name(self.runtime_configuration['module']))
 
         if loglevel:
             self.leverageLogger(loglevel)
diff --git a/intelmq/lib/utils.py b/intelmq/lib/utils.py
index 162ae89eb..78e6ed7ff 100644
--- a/intelmq/lib/utils.py
+++ b/intelmq/lib/utils.py
@@ -852,6 +852,19 @@ def _get_console_entry_points():
     return entries.get("console_scripts", [])  # it's a dict
 
 
+def get_bot_module_name(bot_name: str) -> str:
+    entries = entry_points()
+    if hasattr(entries, "select"):
+        entries = entries.select(name=bot_name, group="console_scripts")
+    else:
+        entries = [entry for entry in entries.get("console_scripts", []) if entry.name == bot_name]
+
+    if not entries:
+        return None
+    else:
+        return entries[0].value.replace(":BOT.run", '')
+
+
 def list_all_bots() -> dict:
     """
     Compile a dictionary with all bots and their parameters.
diff --git a/intelmq/tests/bin/test_intelmqctl.py b/intelmq/tests/bin/test_intelmqctl.py
index f0a594c0e..88ef46a71 100644
--- a/intelmq/tests/bin/test_intelmqctl.py
+++ b/intelmq/tests/bin/test_intelmqctl.py
@@ -125,6 +125,18 @@ def test_check_handles_syntaxerror_when_importing_bots(self):
         self.assertIsNotNone(
             next(filter(lambda l: "SyntaxError in bot 'test-bot'" in l, captured.output)))
 
+    @skip_installation()
+    @mock.patch.object(utils, "get_bot_module_name", mock.Mock(return_value="mocked-module"))
+    def test_check_imports_real_bot_module(self):
+        self._load_default_harmonization()
+        self._extend_config(self.tmp_runtime, self.BOT_CONFIG)
+
+        # raise SyntaxError to stop checking after import
+        with mock.patch.object(ctl.importlib, "import_module", mock.Mock(side_effect=SyntaxError)) as import_mock:
+            self.intelmqctl.check(no_connections=True, check_executables=False)
+
+        import_mock.assert_called_once_with("mocked-module")
+
 
 if __name__ == '__main__':  # pragma: nocover
     unittest.main()
diff --git a/intelmq/tests/lib/test_utils.py b/intelmq/tests/lib/test_utils.py
index f0fac0b8f..48b54032e 100644
--- a/intelmq/tests/lib/test_utils.py
+++ b/intelmq/tests/lib/test_utils.py
@@ -346,6 +346,10 @@ def test_list_all_bots_filters_entrypoints(self):
         )
         self.assertEqual(2, import_mock.call_count)
 
+    def test_get_bot_module_name_builtin_bot(self):
+        found_name = utils.get_bot_module_name("intelmq.bots.collectors.api.collector_api")
+        self.assertEqual("intelmq.bots.collectors.api.collector_api", found_name)
+
     def test_get_bots_settings(self):
         with unittest.mock.patch.object(utils, "get_runtime", new_get_runtime):
             runtime = utils.get_bots_settings()
@@ -381,14 +385,14 @@ def test_load_configuration_yaml(self):
         filename = os.path.join(os.path.dirname(__file__), '../assets/example.yaml')
         self.assertEqual(utils.load_configuration(filename),
                          {
-                            'some_string': 'Hello World!',
-                            'other_string': 'with a : in it',
+            'some_string': 'Hello World!',
+            'other_string': 'with a : in it',
                             'now more': ['values', 'in', 'a', 'list'],
                             'types': -4,
                             'other': True,
                             'final': 0.5,
-                        }
-                        )
+        }
+        )
 
     def test_load_configuration_yaml_invalid(self):
         """ Test load_configuration with an invalid YAML file """