Added generic collector processing functions, a dispatcher (name to function) with unit tests, and moved AWS to collector from regular sysinfo

monkey/infection_monkey/config.py

@@ -125,7 +125,7 @@ class Configuration(object):
 
     finger_classes = []
     exploiter_classes = []
-    system_info_collectors_classes = ["EnvironmentCollector"]
+    system_info_collectors_classes = ["EnvironmentCollector", "AwsCollector"]
 
     # how many victims to look for in a single scan iteration
     victims_max_find = 100

monkey/infection_monkey/system_info/__init__.py

@@ -6,7 +6,6 @@ import psutil
 from enum import IntEnum
 
 from infection_monkey.network.info import get_host_subnets
-from infection_monkey.system_info.aws_collector import AwsCollector
 from infection_monkey.system_info.azure_cred_collector import AzureCollector
 from infection_monkey.system_info.netstat_collector import NetstatCollector
 from system_info.system_info_collectors_handler import SystemInfoCollectorsHandler
@@ -67,7 +66,6 @@ class InfoCollector(object):
         self.get_process_list()
         self.get_network_info()
         self.get_azure_info()
-        self.get_aws_info()
 
         # Collect all plugins
         SystemInfoCollectorsHandler().execute_all_configured()
@@ -155,11 +153,3 @@ class InfoCollector(object):
         except Exception:
             # If we failed to collect azure info, no reason to fail all the collection. Log and continue.
             LOG.error("Failed collecting Azure info.", exc_info=True)
-
-    def get_aws_info(self):
-        # noinspection PyBroadException
-        try:
-            self.info['aws'] = AwsCollector().get_aws_info()
-        except Exception:
-            # If we failed to collect aws info, no reason to fail all the collection. Log and continue.
-            LOG.error("Failed collecting AWS info.", exc_info=True)

monkey/infection_monkey/system_info/collectors/aws_collector.py

@@ -0,0 +1,30 @@
+import logging
+
+from common.cloud.aws.aws_instance import AwsInstance
+from infection_monkey.system_info.system_info_collector import SystemInfoCollector
+
+
+logger = logging.getLogger(__name__)
+
+
+class AwsCollector(SystemInfoCollector):
+    """
+    Extract info from AWS machines.
+    """
+    def __init__(self):
+        super(AwsCollector, self).__init__(name="AwsCollector")
+
+    def collect(self) -> dict:
+        logger.info("Collecting AWS info")
+        aws = AwsInstance()
+        info = {}
+        if aws.is_instance():
+            logger.info("Machine is an AWS instance")
+            info = \
+                {
+                    'instance_id': aws.get_instance_id()
+                }
+        else:
+            logger.info("Machine is NOT an AWS instance")
+
+        return info

monkey/monkey_island/cc/server_config.json

@@ -1,4 +1,4 @@
 {
-  "server_config": "standard",
+  "server_config": "testing",
   "deployment": "develop"
 }

monkey/monkey_island/cc/services/config_schema.py

@@ -111,6 +111,14 @@ SCHEMA = {
                     "title": "Which Environment this machine is on (on prem/cloud)",
                     "attack_techniques": []
                 },
+                {
+                    "type": "string",
+                    "enum": [
+                        "AwsCollector"
+                    ],
+                    "title": "If on AWS, collect more information about the instance",
+                    "attack_techniques": []
+                },
             ],
         },
         "post_breach_acts": {
@@ -455,7 +463,8 @@ SCHEMA = {
                                 "$ref": "#/definitions/system_info_collectors_classes"
                             },
                             "default": [
-                                "EnvironmentCollector"
+                                "EnvironmentCollector",
+                                "AwsCollector"
                             ],
                             "description": "Determines which system information collectors will collect information."
                         },

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

@@ -5,6 +5,7 @@ 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.telemetry.zero_trust_tests.antivirus_existence import test_antivirus_existence
 from monkey_island.cc.services.wmi_handler import WMIHandler
@@ -18,7 +19,7 @@ def process_system_info_telemetry(telemetry_json):
         process_ssh_info,
         process_credential_info,
         process_mimikatz_and_wmi_info,
-        process_aws_data,
+        process_aws_telemetry,
         update_db_with_new_hostname,
         test_antivirus_existence,
         process_environment_telemetry
@@ -116,13 +117,5 @@ def process_mimikatz_and_wmi_info(telemetry_json):
         wmi_handler.process_and_handle_wmi_info()
 
 
-def process_aws_data(telemetry_json):
-    if 'aws' in telemetry_json['data']:
-        if 'instance_id' in telemetry_json['data']['aws']:
-            monkey_id = NodeService.get_monkey_by_guid(telemetry_json['monkey_guid']).get('_id')
-            mongo.db.monkey.update_one({'_id': monkey_id},
-                                       {'$set': {'aws_instance_id': telemetry_json['data']['aws']['instance_id']}})
-
-
 def update_db_with_new_hostname(telemetry_json):
     Monkey.get_single_monkey_by_guid(telemetry_json['monkey_guid']).set_hostname(telemetry_json['data']['hostname'])

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

@@ -0,0 +1,15 @@
+import logging
+
+from monkey_island.cc.models.monkey import Monkey
+
+logger = logging.getLogger(__name__)
+
+
+def process_aws_telemetry(collector_results, monkey_guid):
+    relevant_monkey = Monkey.get_single_monkey_by_guid(monkey_guid)
+
+    if "instance_id" in collector_results:
+        instance_id = collector_results["instance_id"]
+        relevant_monkey.aws_instance_id = instance_id
+        relevant_monkey.save()
+        logger.debug("Updated Monkey {} with aws instance id {}".format(str(relevant_monkey), instance_id))

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

@@ -5,10 +5,8 @@ from monkey_island.cc.models.monkey import Monkey
 logger = logging.getLogger(__name__)
 
 
-def process_environment_telemetry(telemetry_json):
+def process_environment_telemetry(collector_results, monkey_guid):
-    if "EnvironmentCollector" in telemetry_json["data"]["collectors"]:
+    relevant_monkey = Monkey.get_single_monkey_by_guid(monkey_guid)
-        env = telemetry_json["data"]["collectors"]["EnvironmentCollector"]["environment"]
+    relevant_monkey.environment = collector_results
-        relevant_monkey = Monkey.get_single_monkey_by_guid(telemetry_json['monkey_guid'])
-        relevant_monkey.environment = env
     relevant_monkey.save()
-        logger.debug("Updated Monkey {} with env {}".format(str(relevant_monkey), env))
+    logger.debug("Updated Monkey {} with env {}".format(str(relevant_monkey), collector_results))

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

@@ -0,0 +1,36 @@
+import logging
+
+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
+
+logger = logging.getLogger(__name__)
+
+SYSTEM_INFO_COLLECTOR_TO_TELEMETRY_PROCESSOR = {
+    "AwsCollector": process_aws_telemetry,
+    "EnvironmentCollector": process_environment_telemetry,
+}
+
+
+class SystemInfoTelemetryDispatcher(object):
+    def __init__(self, collector_to_parsing_function=None):
+        if collector_to_parsing_function is None:
+            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):
+        if "collectors" in telemetry_json["data"]:
+            self.send_each_result_to_relevant_processor(telemetry_json)
+
+    def send_each_result_to_relevant_processor(self, telemetry_json):
+        relevant_monkey_guid = telemetry_json['monkey_guid']
+        for collector_name, collector_results in telemetry_json["data"]["collectors"].items():
+            if collector_name in self.collector_to_parsing_function:
+                # noinspection PyBroadException
+                try:
+                    self.collector_to_parsing_function[collector_name](collector_results, relevant_monkey_guid)
+                except Exception as e:
+                    logger.error(
+                        "Error {} while processing {} system info telemetry".format(str(e), collector_name),
+                        exc_info=True)
+            else:
+                logger.warning("Unknown system info collector name: {}".format(collector_name))

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

@@ -0,0 +1,65 @@
+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
+
+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()
+
+        dispatcher = SystemInfoTelemetryDispatcher(TEST_SYS_INFO_TO_PROCESSING)
+
+        # Bad format telem JSONs - throws
+        bad_empty_telem_json = {}
+        self.assertRaises(KeyError, dispatcher.dispatch_to_relevant_collector, 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)
+        bad_no_monkey_telem_json = {"data": {"collectors": {"AwsCollector": "Bla"}}}
+        self.assertRaises(KeyError, dispatcher.dispatch_to_relevant_collector, 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)
+
+    def test_dispatch_to_relevant_collector(self):
+        self.fail_if_not_testing_env()
+        self.clean_monkey_db()
+
+        a_monkey = Monkey(guid=str(uuid.uuid4()))
+        a_monkey.save()
+
+        dispatcher = SystemInfoTelemetryDispatcher()
+
+        # JSON with results - make sure functions are called
+        instance_id = "i-0bd2c14bd4c7d703f"
+        telem_json = {
+            "data": {
+                "collectors": {
+                    "AwsCollector": {"instance_id": instance_id},
+                }
+            },
+            "monkey_guid": a_monkey.guid
+        }
+        dispatcher.dispatch_to_relevant_collector(telem_json)
+
+        self.assertEquals(Monkey.get_single_monkey_by_guid(a_monkey.guid).aws_instance_id, instance_id)