Agent: Refactor AWS collector

This commit is contained in:
Ilija Lazoroski 2022-02-11 12:22:41 +01:00
parent 5a8c072d6a
commit 1f76a42279
4 changed files with 45 additions and 0 deletions

View File

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

View File

@ -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")

View File

@ -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)

View File

@ -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")