Using the dispatcher instead of naming the functions one by one + optimize imports

monkey/monkey_island/cc/services/telemetry/processing/system_info.py

@@ -1,28 +1,27 @@
 import logging
 
-from monkey_island.cc.database import mongo
+from monkey_island.cc.encryptor import encryptor
 from monkey_island.cc.models import Monkey
 from monkey_island.cc.services import mimikatz_utils
-from monkey_island.cc.services.node import NodeService
 from monkey_island.cc.services.config import ConfigService
-from monkey_island.cc.services.telemetry.processing.system_info_collectors.aws import process_aws_telemetry
-from monkey_island.cc.services.telemetry.processing.system_info_collectors.environment import process_environment_telemetry
+from monkey_island.cc.services.node import NodeService
+from monkey_island.cc.services.telemetry.processing.system_info_collectors.system_info_telemetry_dispatcher import \
+    SystemInfoTelemetryDispatcher
 from monkey_island.cc.services.telemetry.zero_trust_tests.antivirus_existence import test_antivirus_existence
 from monkey_island.cc.services.wmi_handler import WMIHandler
-from monkey_island.cc.encryptor import encryptor
 
 logger = logging.getLogger(__name__)
 
 
 def process_system_info_telemetry(telemetry_json):
+    dispatcher = SystemInfoTelemetryDispatcher()
     telemetry_processing_stages = [
         process_ssh_info,
         process_credential_info,
         process_mimikatz_and_wmi_info,
-        process_aws_telemetry,
         update_db_with_new_hostname,
         test_antivirus_existence,
-        process_environment_telemetry
+        dispatcher.dispatch_to_relevant_collectors
     ]
 
     # Calling safe_process_telemetry so if one of the stages fail, we log and move on instead of failing the rest of

monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py

@@ -17,7 +17,7 @@ class SystemInfoTelemetryDispatcher(object):
             collector_to_parsing_function = SYSTEM_INFO_COLLECTOR_TO_TELEMETRY_PROCESSOR
         self.collector_to_parsing_function = collector_to_parsing_function
 
-    def dispatch_to_relevant_collector(self, telemetry_json):
+    def dispatch_to_relevant_collectors(self, telemetry_json):
         if "collectors" in telemetry_json["data"]:
             self.send_each_result_to_relevant_processor(telemetry_json)
 

monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py

@@ -1,25 +1,17 @@
-from importlib import reload
-from unittest import mock
-from unittest.mock import MagicMock
-
 import uuid
 
 from monkey_island.cc.models import Monkey
 from monkey_island.cc.services.telemetry.processing.system_info_collectors.system_info_telemetry_dispatcher import \
     SystemInfoTelemetryDispatcher
-from monkey_island.cc.testing.IslandTestCase import IslandTestCase
 from monkey_island.cc.services.telemetry.processing.system_info_collectors.system_info_telemetry_dispatcher import \
     process_aws_telemetry
+from monkey_island.cc.testing.IslandTestCase import IslandTestCase
 
 TEST_SYS_INFO_TO_PROCESSING = {
     "AwsCollector": process_aws_telemetry,
 }
 
 
-def do_nothing(x, y):
-    pass
-
-
 class SystemInfoTelemetryDispatcherTest(IslandTestCase):
     def test_dispatch_to_relevant_collector_bad_inputs(self):
         self.fail_if_not_testing_env()
@@ -28,18 +20,18 @@ class SystemInfoTelemetryDispatcherTest(IslandTestCase):
 
         # Bad format telem JSONs - throws
         bad_empty_telem_json = {}
-        self.assertRaises(KeyError, dispatcher.dispatch_to_relevant_collector, bad_empty_telem_json)
+        self.assertRaises(KeyError, dispatcher.dispatch_to_relevant_collectors, bad_empty_telem_json)
         bad_no_data_telem_json = {"monkey_guid": "bla"}
-        self.assertRaises(KeyError, dispatcher.dispatch_to_relevant_collector, bad_no_data_telem_json)
+        self.assertRaises(KeyError, dispatcher.dispatch_to_relevant_collectors, bad_no_data_telem_json)
         bad_no_monkey_telem_json = {"data": {"collectors": {"AwsCollector": "Bla"}}}
-        self.assertRaises(KeyError, dispatcher.dispatch_to_relevant_collector, bad_no_monkey_telem_json)
+        self.assertRaises(KeyError, dispatcher.dispatch_to_relevant_collectors, bad_no_monkey_telem_json)
 
         # Telem JSON with no collectors - nothing gets dispatched
         good_telem_no_collectors = {"monkey_guid": "bla", "data": {"bla": "bla"}}
         good_telem_empty_collectors = {"monkey_guid": "bla", "data": {"bla": "bla", "collectors": {}}}
 
-        dispatcher.dispatch_to_relevant_collector(good_telem_no_collectors)
-        dispatcher.dispatch_to_relevant_collector(good_telem_empty_collectors)
+        dispatcher.dispatch_to_relevant_collectors(good_telem_no_collectors)
+        dispatcher.dispatch_to_relevant_collectors(good_telem_empty_collectors)
 
     def test_dispatch_to_relevant_collector(self):
         self.fail_if_not_testing_env()
@@ -60,6 +52,6 @@ class SystemInfoTelemetryDispatcherTest(IslandTestCase):
             },
             "monkey_guid": a_monkey.guid
         }
-        dispatcher.dispatch_to_relevant_collector(telem_json)
+        dispatcher.dispatch_to_relevant_collectors(telem_json)
 
         self.assertEquals(Monkey.get_single_monkey_by_guid(a_monkey.guid).aws_instance_id, instance_id)