diff --git a/monkey/common/cloud/aws/aws_instance.py b/monkey/common/cloud/aws/aws_instance.py index 4bdc89bf3..76f4ab258 100644 --- a/monkey/common/cloud/aws/aws_instance.py +++ b/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: diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index b4a8a8566..c0b5d17b3 100644 --- a/monkey/infection_monkey/monkey.py +++ b/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, diff --git a/monkey/infection_monkey/telemetry/aws_instance_telem.py b/monkey/infection_monkey/telemetry/aws_instance_telem.py index d421d6d00..d2469b971 100644 --- a/monkey/infection_monkey/telemetry/aws_instance_telem.py +++ b/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): - def __init__(self, aws_instance_info): +class AWSInstanceTelemetry(BaseTelem): + 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) diff --git a/monkey/infection_monkey/utils/aws_environment_check.py b/monkey/infection_monkey/utils/aws_environment_check.py index 03ee5a579..31ff40186 100644 --- a/monkey/infection_monkey/utils/aws_environment_check.py +++ b/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()