From 6815433a85d493778f085c908cbd79f1d4c58f0b Mon Sep 17 00:00:00 2001 From: Shay Nehmad Date: Sun, 19 Jan 2020 21:39:36 +0200 Subject: [PATCH] Using the dispatcher instead of naming the functions one by one + optimize imports --- .../telemetry/processing/system_info.py | 13 +++++------ .../system_info_telemetry_dispatcher.py | 2 +- .../test_system_info_telemetry_dispatcher.py | 22 ++++++------------- 3 files changed, 14 insertions(+), 23 deletions(-) diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info.py index 915fa7a25..6734c6725 100644 --- a/monkey/monkey_island/cc/services/telemetry/processing/system_info.py +++ b/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 diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py index 661034efb..64fb146ab 100644 --- a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py +++ b/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) diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py index db36cd5bb..4db5352c8 100644 --- a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py +++ b/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)