Merge pull request #1712 from guardicore/1696-refactor-aws-collector

Agent: Refactor AWS collector

monkey/common/cloud/aws/aws_instance.py

@@ -12,6 +12,8 @@ ACCOUNT_ID_KEY = "accountId"
 
 logger = logging.getLogger(__name__)
 
+AWS_TIMEOUT = 2
+
 
 class AwsInstance(CloudInstance):
     """
@@ -28,12 +30,14 @@ class AwsInstance(CloudInstance):
 
         try:
             response = requests.get(
-                AWS_LATEST_METADATA_URI_PREFIX + "meta-data/instance-id", timeout=2
+                AWS_LATEST_METADATA_URI_PREFIX + "meta-data/instance-id",
+                timeout=AWS_TIMEOUT,
             )
             self.instance_id = response.text if response else None
             self.region = self._parse_region(
                 requests.get(
-                    AWS_LATEST_METADATA_URI_PREFIX + "meta-data/placement/availability-zone"
+                    AWS_LATEST_METADATA_URI_PREFIX + "meta-data/placement/availability-zone",
+                    timeout=AWS_TIMEOUT,
                 ).text
             )
         except (requests.RequestException, IOError) as e:
@@ -42,7 +46,8 @@ class AwsInstance(CloudInstance):
         try:
             self.account_id = self._extract_account_id(
                 requests.get(
-                    AWS_LATEST_METADATA_URI_PREFIX + "dynamic/instance-identity/document", timeout=2
+                    AWS_LATEST_METADATA_URI_PREFIX + "dynamic/instance-identity/document",
+                    timeout=AWS_TIMEOUT,
                 ).text
             )
         except (requests.RequestException, json.decoder.JSONDecodeError, IOError) as e:

monkey/common/common_consts/system_info_collectors_names.py

@@ -1,3 +1,2 @@
-AWS_COLLECTOR = "AwsCollector"
 PROCESS_LIST_COLLECTOR = "ProcessListCollector"
 MIMIKATZ_COLLECTOR = "MimikatzCollector"

monkey/common/common_consts/telem_categories.py

@@ -8,3 +8,4 @@ class TelemCategoryEnum:
     TUNNEL = "tunnel"
     ATTACK = "attack"
     FILE_ENCRYPTION = "file_encryption"
+    AWS_INFO = "aws_info"

monkey/infection_monkey/monkey.py

@@ -34,6 +34,7 @@ from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter im
 )
 from infection_monkey.telemetry.state_telem import StateTelem
 from infection_monkey.telemetry.tunnel_telem import TunnelTelem
+from infection_monkey.utils.aws_environment_check import run_aws_environment_check
 from infection_monkey.utils.environment import is_windows_os
 from infection_monkey.utils.monkey_dir import get_monkey_dir_path, remove_monkey_dir
 from infection_monkey.utils.monkey_log_path import get_monkey_log_path
@@ -52,6 +53,7 @@ class InfectionMonkey:
         self._default_server = self._opts.server
         # TODO used in propogation phase
         self._monkey_inbound_tunnel = None
+        self.telemetry_messenger = LegacyTelemetryMessengerAdapter()
 
     @staticmethod
     def _get_arguments(args):
@@ -85,6 +87,8 @@ class InfectionMonkey:
         if is_windows_os():
             T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
 
+        run_aws_environment_check(self.telemetry_messenger)
+
         should_stop = ControlChannel(WormConfiguration.current_server, GUID).should_agent_stop()
         if should_stop:
             logger.info("The Monkey Island has instructed this agent to stop")
@@ -171,7 +175,7 @@ class InfectionMonkey:
 
         self._master = AutomatedMaster(
             puppet,
-            LegacyTelemetryMessengerAdapter(),
+            self.telemetry_messenger,
             victim_host_factory,
             ControlChannel(self._default_server, GUID),
             local_network_interfaces,

monkey/infection_monkey/system_info/collectors/aws_collector.py

@@ -1,28 +0,0 @@
-import logging
-
-from common.cloud.aws.aws_instance import AwsInstance
-from common.common_consts.system_info_collectors_names import AWS_COLLECTOR
-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().__init__(name=AWS_COLLECTOR)
-
-    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/infection_monkey/telemetry/aws_instance_telem.py

@@ -0,0 +1,18 @@
+from common.common_consts.telem_categories import TelemCategoryEnum
+from infection_monkey.telemetry.base_telem import BaseTelem
+
+
+class AWSInstanceTelemetry(BaseTelem):
+    def __init__(self, aws_instance_id: str):
+        """
+        Default AWS instance telemetry constructor
+        """
+        self.aws_instance_info = {"instance_id": aws_instance_id}
+
+    telem_category = TelemCategoryEnum.AWS_INFO
+
+    def get_data(self):
+        return self.aws_instance_info
+
+    def send(self, log_data=False):
+        super(AWSInstanceTelemetry, self).send(log_data)

monkey/infection_monkey/utils/aws_environment_check.py

@@ -0,0 +1,34 @@
+import logging
+
+from common.cloud.aws.aws_instance import AwsInstance
+from infection_monkey.telemetry.aws_instance_telem import AWSInstanceTelemetry
+from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter import (
+    LegacyTelemetryMessengerAdapter,
+)
+from infection_monkey.utils.threading import create_daemon_thread
+
+logger = logging.getLogger(__name__)
+
+
+def _running_on_aws(aws_instance: AwsInstance) -> bool:
+    return aws_instance.is_instance()
+
+
+def _report_aws_environment(telemetry_messenger: LegacyTelemetryMessengerAdapter):
+    logger.info("Collecting AWS info")
+
+    aws_instance = AwsInstance()
+
+    if _running_on_aws(aws_instance):
+        logger.info("Machine is an AWS instance")
+        telemetry_messenger.send_telemetry(AWSInstanceTelemetry(aws_instance.get_instance_id()))
+    else:
+        logger.info("Machine is NOT an AWS instance")
+
+
+def run_aws_environment_check(telemetry_messenger: LegacyTelemetryMessengerAdapter):
+    logger.info("AWS environment check initiated.")
+    aws_environment_thread = create_daemon_thread(
+        target=_report_aws_environment, args=(telemetry_messenger,)
+    )
+    aws_environment_thread.start()

monkey/monkey_island/cc/services/config_schema/definitions/system_info_collector_classes.py

@@ -1,5 +1,4 @@
 from common.common_consts.system_info_collectors_names import (
-    AWS_COLLECTOR,
     MIMIKATZ_COLLECTOR,
     PROCESS_LIST_COLLECTOR,
 )
@@ -17,15 +16,6 @@ SYSTEM_INFO_COLLECTOR_CLASSES = {
             "info": "Collects credentials from Windows credential manager.",
             "attack_techniques": ["T1003", "T1005"],
         },
-        {
-            "type": "string",
-            "enum": [AWS_COLLECTOR],
-            "title": "AWS Collector",
-            "safe": True,
-            "info": "If on AWS, collects more information about the AWS instance "
-            "currently running on.",
-            "attack_techniques": ["T1082"],
-        },
         {
             "type": "string",
             "enum": [PROCESS_LIST_COLLECTOR],

monkey/monkey_island/cc/services/config_schema/monkey.py

@@ -1,5 +1,4 @@
 from common.common_consts.system_info_collectors_names import (
-    AWS_COLLECTOR,
     MIMIKATZ_COLLECTOR,
     PROCESS_LIST_COLLECTOR,
 )
@@ -86,7 +85,6 @@ MONKEY = {
                     "uniqueItems": True,
                     "items": {"$ref": "#/definitions/system_info_collector_classes"},
                     "default": [
-                        AWS_COLLECTOR,
                         PROCESS_LIST_COLLECTOR,
                         MIMIKATZ_COLLECTOR,
                     ],

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

@@ -5,11 +5,11 @@ from monkey_island.cc.models.monkey import Monkey
 logger = logging.getLogger(__name__)
 
 
-def process_aws_telemetry(collector_results, monkey_guid):
+def process_aws_telemetry(telemetry_json):
-    relevant_monkey = Monkey.get_single_monkey_by_guid(monkey_guid)
+    relevant_monkey = Monkey.get_single_monkey_by_guid(telemetry_json["monkey_guid"])
 
-    if "instance_id" in collector_results:
+    if "instance_id" in telemetry_json["data"]:
-        instance_id = collector_results["instance_id"]
+        instance_id = telemetry_json["data"]["instance_id"]
         relevant_monkey.aws_instance_id = instance_id
         relevant_monkey.save()
         logger.debug(

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

@@ -1,6 +1,7 @@
 import logging
 
 from common.common_consts.telem_categories import TelemCategoryEnum
+from monkey_island.cc.services.telemetry.processing.aws_info import process_aws_telemetry
 from monkey_island.cc.services.telemetry.processing.exploit import process_exploit_telemetry
 from monkey_island.cc.services.telemetry.processing.post_breach import process_post_breach_telemetry
 from monkey_island.cc.services.telemetry.processing.scan import process_scan_telemetry
@@ -17,6 +18,7 @@ TELEMETRY_CATEGORY_TO_PROCESSING_FUNC = {
     TelemCategoryEnum.SCAN: process_scan_telemetry,
     TelemCategoryEnum.SYSTEM_INFO: process_system_info_telemetry,
     TelemCategoryEnum.POST_BREACH: process_post_breach_telemetry,
+    TelemCategoryEnum.AWS_INFO: process_aws_telemetry,
     # `lambda *args, **kwargs: None` is a no-op.
     TelemCategoryEnum.TRACE: lambda *args, **kwargs: None,
     TelemCategoryEnum.ATTACK: lambda *args, **kwargs: None,

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

@@ -1,10 +1,7 @@
 import logging
 import typing
 
-from common.common_consts.system_info_collectors_names import AWS_COLLECTOR, PROCESS_LIST_COLLECTOR
+from common.common_consts.system_info_collectors_names import PROCESS_LIST_COLLECTOR
-from monkey_island.cc.services.telemetry.processing.system_info_collectors.aws import (
-    process_aws_telemetry,
-)
 from monkey_island.cc.services.telemetry.zero_trust_checks.antivirus_existence import (
     check_antivirus_existence,
 )
@@ -12,7 +9,6 @@ from monkey_island.cc.services.telemetry.zero_trust_checks.antivirus_existence i
 logger = logging.getLogger(__name__)
 
 SYSTEM_INFO_COLLECTOR_TO_TELEMETRY_PROCESSORS = {
-    AWS_COLLECTOR: [process_aws_telemetry],
     PROCESS_LIST_COLLECTOR: [check_antivirus_existence],
 }
 

monkey/tests/data_for_tests/monkey_configs/automated_master_config.json

@@ -104,7 +104,6 @@
             }
         },
         "system_info_collector_classes": [
-            "AwsCollector",
             "ProcessListCollector",
             "MimikatzCollector"
         ]

monkey/tests/data_for_tests/monkey_configs/flat_config.json

@@ -101,7 +101,6 @@
     "smb_service_name": "InfectionMonkey",
     "subnet_scan_list": ["192.168.1.50", "192.168.56.0/24", "10.0.33.0/30"],
     "system_info_collector_classes": [
-        "AwsCollector",
         "ProcessListCollector",
         "MimikatzCollector"
     ],

monkey/tests/data_for_tests/monkey_configs/monkey_config_standard.json

@@ -147,7 +147,6 @@
       "system_info": {
         "system_info_collector_classes": [
           "environmentcollector",
-          "awscollector",
           "hostnamecollector",
           "processlistcollector",
           "mimikatzcollector"

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

@@ -1,61 +0,0 @@
-import uuid
-
-import pytest
-
-from monkey_island.cc.models import Monkey
-from monkey_island.cc.services.telemetry.processing.system_info_collectors.system_info_telemetry_dispatcher import (  # noqa: E501
-    SystemInfoTelemetryDispatcher,
-    process_aws_telemetry,
-)
-
-TEST_SYS_INFO_TO_PROCESSING = {
-    "AwsCollector": [process_aws_telemetry],
-}
-
-
-class TestSystemInfoTelemetryDispatcher:
-    def test_dispatch_to_relevant_collector_bad_inputs(self):
-        dispatcher = SystemInfoTelemetryDispatcher(TEST_SYS_INFO_TO_PROCESSING)
-
-        # Bad format telem JSONs - throws
-        bad_empty_telem_json = {}
-        with pytest.raises(KeyError):
-            dispatcher.dispatch_collector_results_to_relevant_processors(bad_empty_telem_json)
-
-        bad_no_data_telem_json = {"monkey_guid": "bla"}
-        with pytest.raises(KeyError):
-            dispatcher.dispatch_collector_results_to_relevant_processors(bad_no_data_telem_json)
-
-        bad_no_monkey_telem_json = {"data": {"collectors": {"AwsCollector": "Bla"}}}
-        with pytest.raises(KeyError):
-            dispatcher.dispatch_collector_results_to_relevant_processors(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_collector_results_to_relevant_processors(good_telem_no_collectors)
-        dispatcher.dispatch_collector_results_to_relevant_processors(good_telem_empty_collectors)
-
-    def test_dispatch_to_relevant_collector(self):
-        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_collector_results_to_relevant_processors(telem_json)
-
-        assert Monkey.get_single_monkey_by_guid(a_monkey.guid).aws_instance_id == instance_id

vulture_allowlist.py

@@ -96,7 +96,6 @@ AccountDiscovery  # unused class (monkey/infection_monkey/post_breach/actions/di
 ModifyShellStartupFiles  # unused class (monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py:11)
 Timestomping  # unused class (monkey/infection_monkey/post_breach/actions/timestomping.py:6)
 SignedScriptProxyExecution  # unused class (monkey/infection_monkey/post_breach/actions/use_signed_scripts.py:15)
-AwsCollector  # unused class (monkey/infection_monkey/system_info/collectors/aws_collector.py:15)
 EnvironmentCollector  # unused class (monkey/infection_monkey/system_info/collectors/environment_collector.py:19)
 HostnameCollector  # unused class (monkey/infection_monkey/system_info/collectors/hostname_collector.py:10)
 ProcessListCollector  # unused class (monkey/infection_monkey/system_info/collectors/process_list_collector.py:18)