Agent: Run AWS Environment check in a thread

* Use Telemetry Messenger to send AWS telemetry
* Send only instance_id to AWS Instance Telemetry
* Rename AwsInstanceTelemetry to AWSInstanceTelemetry

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/infection_monkey/monkey.py

@@ -34,7 +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 report_aws_environment
+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
@@ -53,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):
@@ -86,7 +87,7 @@ class InfectionMonkey:
         if is_windows_os():
             T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
 
-        report_aws_environment()
+        run_aws_environment_check(self.telemetry_messenger)
 
         should_stop = ControlChannel(WormConfiguration.current_server, GUID).should_agent_stop()
         if should_stop:
@@ -174,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/telemetry/aws_instance_telem.py

@@ -2,13 +2,12 @@ from common.common_consts.telem_categories import TelemCategoryEnum
 from infection_monkey.telemetry.base_telem import BaseTelem
 
 
-class AwsInstanceTelemetry(BaseTelem):
+class AWSInstanceTelemetry(BaseTelem):
-    def __init__(self, aws_instance_info):
+    def __init__(self, aws_instance_id: str):
         """
         Default AWS instance telemetry constructor
-        :param aws_instance_info: Aws Instance info
         """
-        self.aws_instance_info = aws_instance_info
+        self.aws_instance_info = {"instance_id": aws_instance_id}
 
     telem_category = TelemCategoryEnum.AWS_INFO
 
@@ -16,4 +15,4 @@ class AwsInstanceTelemetry(BaseTelem):
         return self.aws_instance_info
 
     def send(self, log_data=False):
-        super(AwsInstanceTelemetry, self).send(log_data)
+        super(AWSInstanceTelemetry, self).send(log_data)

monkey/infection_monkey/utils/aws_environment_check.py

@@ -1,7 +1,11 @@
 import logging
 
 from common.cloud.aws.aws_instance import AwsInstance
-from infection_monkey.telemetry.aws_instance_telem import AwsInstanceTelemetry
+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__)
 
@@ -10,13 +14,21 @@ def _running_on_aws(aws_instance: AwsInstance) -> bool:
     return aws_instance.is_instance()
 
 
-def report_aws_environment():
+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")
-        AwsInstanceTelemetry({"instance_id": aws_instance.get_instance_id()}).send()
+        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()