From 1f76a422795cced6110f3f0d017855d155c6edf4 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Fri, 11 Feb 2022 12:22:41 +0100
Subject: [PATCH] Agent: Refactor AWS collector

---
 .../common/common_consts/telem_categories.py  |  1 +
 monkey/infection_monkey/monkey.py             |  3 +++
 .../telemetry/aws_instance_telem.py           | 19 ++++++++++++++++
 .../utils/aws_environment_check.py            | 22 +++++++++++++++++++
 4 files changed, 45 insertions(+)
 create mode 100644 monkey/infection_monkey/telemetry/aws_instance_telem.py
 create mode 100644 monkey/infection_monkey/utils/aws_environment_check.py

diff --git a/monkey/common/common_consts/telem_categories.py b/monkey/common/common_consts/telem_categories.py
index dc6524c7b..c9d3f82bd 100644
--- a/monkey/common/common_consts/telem_categories.py
+++ b/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"
diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index e06a39689..b4a8a8566 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/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 report_aws_environment
 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
@@ -85,6 +86,8 @@ class InfectionMonkey:
         if is_windows_os():
             T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
 
+        report_aws_environment()
+
         should_stop = ControlChannel(WormConfiguration.current_server, GUID).should_agent_stop()
         if should_stop:
             logger.info("The Monkey Island has instructed this agent to stop")
diff --git a/monkey/infection_monkey/telemetry/aws_instance_telem.py b/monkey/infection_monkey/telemetry/aws_instance_telem.py
new file mode 100644
index 000000000..d421d6d00
--- /dev/null
+++ b/monkey/infection_monkey/telemetry/aws_instance_telem.py
@@ -0,0 +1,19 @@
+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):
+        """
+        Default AWS instance telemetry constructor
+        :param aws_instance_info: Aws Instance info
+        """
+        self.aws_instance_info = aws_instance_info
+
+    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)
diff --git a/monkey/infection_monkey/utils/aws_environment_check.py b/monkey/infection_monkey/utils/aws_environment_check.py
new file mode 100644
index 000000000..03ee5a579
--- /dev/null
+++ b/monkey/infection_monkey/utils/aws_environment_check.py
@@ -0,0 +1,22 @@
+import logging
+
+from common.cloud.aws.aws_instance import AwsInstance
+from infection_monkey.telemetry.aws_instance_telem import AwsInstanceTelemetry
+
+logger = logging.getLogger(__name__)
+
+
+def _running_on_aws(aws_instance: AwsInstance) -> bool:
+    return aws_instance.is_instance()
+
+
+def report_aws_environment():
+    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()
+    else:
+        logger.info("Machine is NOT an AWS instance")