From e83848c8a4c8e709069fbf23622050555617077d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 07:09:11 -0400 Subject: [PATCH 1/7] Island: Add AWSInstance to the DIContainer --- monkey/monkey_island/cc/services/initialize.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index faa3bcef9..bcc11676d 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -2,6 +2,7 @@ from pathlib import Path from threading import Thread from common import DIContainer +from common.aws import AWSInstance from monkey_island.cc.services import DirectoryFileStorageService, IFileStorageService, aws_service from monkey_island.cc.services.post_breach_files import PostBreachFilesService from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService @@ -14,6 +15,7 @@ def initialize_services(data_dir: Path) -> DIContainer: container.register_instance( IFileStorageService, DirectoryFileStorageService(data_dir / "custom_pbas") ) + container.register_instance(AWSInstance, AWSInstance()) # Takes a while so it's best to start it in the background Thread(target=aws_service.initialize, name="AwsService initialization", daemon=True).start() From 0f4b69a6f7a1b7479bd4bf99577ea0b6e9a7bb78 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 08:27:10 -0400 Subject: [PATCH 2/7] Island: Add stateful AWSService --- monkey/monkey_island/cc/services/__init__.py | 2 + .../monkey_island/cc/services/aws_service.py | 12 ++++ .../cc/services/test_aws_service.py | 56 +++++++++++++++++++ 3 files changed, 70 insertions(+) diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index 43aa39382..7d96ee5c4 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/monkey/monkey_island/cc/services/__init__.py @@ -3,3 +3,5 @@ from .directory_file_storage_service import DirectoryFileStorageService from .authentication.authentication_service import AuthenticationService from .authentication.json_file_user_datastore import JsonFileUserDatastore + +from .aws_service import AWSService diff --git a/monkey/monkey_island/cc/services/aws_service.py b/monkey/monkey_island/cc/services/aws_service.py index 12432b8b7..fe7f51487 100644 --- a/monkey/monkey_island/cc/services/aws_service.py +++ b/monkey/monkey_island/cc/services/aws_service.py @@ -15,6 +15,18 @@ IP_ADDRESS_KEY = "IPAddress" logger = logging.getLogger(__name__) +class AWSService: + def __init__(self, aws_instance: AWSInstance): + self._aws_instance = aws_instance + + def island_is_running_on_aws(self) -> bool: + return self._aws_instance.is_instance + + @property + def island_aws_instance(self) -> AWSInstance: + return self._aws_instance + + def filter_instance_data_from_aws_response(response): return [ { diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_aws_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_aws_service.py index 0d8a71f36..bd85595f5 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_aws_service.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_aws_service.py @@ -1,6 +1,12 @@ import json +import threading +from typing import Optional from unittest import TestCase +import pytest + +from common.aws import AWSInstance +from monkey_island.cc.services import AWSService from monkey_island.cc.services.aws_service import filter_instance_data_from_aws_response @@ -54,3 +60,53 @@ class TestAwsService(TestCase): filter_instance_data_from_aws_response(json.loads(json_response_full)), [{"instance_id": "string", "ip_address": "string", "name": "string", "os": "string"}], ) + + +class StubAWSInstance(AWSInstance): + def __init__( + self, + instance_id: Optional[str] = None, + region: Optional[str] = None, + account_id: Optional[str] = None, + ): + self._instance_id = instance_id + self._region = region + self._account_id = account_id + + self._initialization_complete = threading.Event() + self._initialization_complete.set() + + +def test_aws_is_on_aws__true(): + aws_instance = StubAWSInstance("1") + aws_service = AWSService(aws_instance) + assert aws_service.island_is_running_on_aws() is True + + +def test_aws_is_on_aws__False(): + aws_instance = StubAWSInstance() + aws_service = AWSService(aws_instance) + assert aws_service.island_is_running_on_aws() is False + + +INSTANCE_ID = "1" +REGION = "2" +ACCOUNT_ID = "3" + + +@pytest.fixture +def aws_service(): + aws_instance = StubAWSInstance(INSTANCE_ID, REGION, ACCOUNT_ID) + return AWSService(aws_instance) + + +def test_instance_id(aws_service): + assert aws_service.island_aws_instance.instance_id == INSTANCE_ID + + +def test_region(aws_service): + assert aws_service.island_aws_instance.region == REGION + + +def test_account_id(aws_service): + assert aws_service.island_aws_instance.account_id == ACCOUNT_ID From acabc835d458e55ca10ec95b7cc6f2f468996368 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 09:22:33 -0400 Subject: [PATCH 3/7] Island: Add run_agent_on_managed_instances() to AWSService --- monkey/monkey_island/cc/services/aws_service.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/services/aws_service.py b/monkey/monkey_island/cc/services/aws_service.py index fe7f51487..1c15bb6d5 100644 --- a/monkey/monkey_island/cc/services/aws_service.py +++ b/monkey/monkey_island/cc/services/aws_service.py @@ -1,5 +1,5 @@ import logging -from typing import Optional +from typing import Iterable, Optional import boto3 import botocore @@ -26,6 +26,13 @@ class AWSService: def island_aws_instance(self) -> AWSInstance: return self._aws_instance + def run_agent_on_managed_instances(self, instance_ids: Iterable[str]): + for id_ in instance_ids: + self._run_agent_on_managed_instance(id_) + + def _run_agent_on_managed_instance(self, instance_id: str): + pass + def filter_instance_data_from_aws_response(response): return [ From 8995eb5d2f92121de66cb3ddf3fd2cc4d0ca7594 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 10:09:17 -0400 Subject: [PATCH 4/7] Island: Add get_managed_instances() to AWSService --- .../monkey_island/cc/services/aws_service.py | 83 ++++------ .../cc/services/test_aws_service.py | 150 +++++++++++------- 2 files changed, 124 insertions(+), 109 deletions(-) diff --git a/monkey/monkey_island/cc/services/aws_service.py b/monkey/monkey_island/cc/services/aws_service.py index 1c15bb6d5..d589fad96 100644 --- a/monkey/monkey_island/cc/services/aws_service.py +++ b/monkey/monkey_island/cc/services/aws_service.py @@ -1,5 +1,5 @@ import logging -from typing import Iterable, Optional +from typing import Any, Dict, Iterable, Sequence import boto3 import botocore @@ -26,6 +26,29 @@ class AWSService: def island_aws_instance(self) -> AWSInstance: return self._aws_instance + def get_managed_instances(self) -> Sequence[Dict[str, str]]: + raw_managed_instances_info = self._get_raw_managed_instances() + return _filter_instance_info_from_aws_response(raw_managed_instances_info) + + def _get_raw_managed_instances(self) -> Sequence[Dict[str, Any]]: + """ + Get the information for all instances with the relevant roles. + + This function will assume that it's running on an EC2 instance with the correct IAM role. + See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html#iam + -role for details. + + :raises: botocore.exceptions.ClientError if can't describe local instance information. + :return: All visible instances from this instance + """ + local_ssm_client = boto3.client("ssm", self.island_aws_instance.region) + try: + response = local_ssm_client.describe_instance_information() + return response[INSTANCE_INFORMATION_LIST_KEY] + except botocore.exceptions.ClientError as err: + logger.warning("AWS client error while trying to get manage dinstances: {err}") + raise err + def run_agent_on_managed_instances(self, instance_ids: Iterable[str]): for id_ in instance_ids: self._run_agent_on_managed_instance(id_) @@ -34,59 +57,13 @@ class AWSService: pass -def filter_instance_data_from_aws_response(response): +def _filter_instance_info_from_aws_response(raw_managed_instances_info: Sequence[Dict[str, Any]]): return [ { - "instance_id": x[INSTANCE_ID_KEY], - "name": x[COMPUTER_NAME_KEY], - "os": x[PLATFORM_TYPE_KEY].lower(), - "ip_address": x[IP_ADDRESS_KEY], + "instance_id": managed_instance[INSTANCE_ID_KEY], + "name": managed_instance[COMPUTER_NAME_KEY], + "os": managed_instance[PLATFORM_TYPE_KEY].lower(), + "ip_address": managed_instance[IP_ADDRESS_KEY], } - for x in response[INSTANCE_INFORMATION_LIST_KEY] + for managed_instance in raw_managed_instances_info ] - - -aws_instance: Optional[AWSInstance] = None - - -def initialize(): - global aws_instance - aws_instance = AWSInstance() - - -def is_on_aws(): - return aws_instance.is_instance - - -def get_region(): - return aws_instance.region - - -def get_account_id(): - return aws_instance.account_id - - -def get_client(client_type): - return boto3.client(client_type, region_name=aws_instance.region) - - -def get_instances(): - """ - Get the information for all instances with the relevant roles. - - This function will assume that it's running on an EC2 instance with the correct IAM role. - See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html#iam - -role for details. - - :raises: botocore.exceptions.ClientError if can't describe local instance information. - :return: All visible instances from this instance - """ - local_ssm_client = boto3.client("ssm", aws_instance.region) - try: - response = local_ssm_client.describe_instance_information() - - filtered_instances_data = filter_instance_data_from_aws_response(response) - return filtered_instances_data - except botocore.exceptions.ClientError as e: - logger.warning("AWS client error while trying to get instances: " + e) - raise e diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_aws_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_aws_service.py index bd85595f5..8be306310 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_aws_service.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_aws_service.py @@ -1,65 +1,75 @@ -import json import threading -from typing import Optional -from unittest import TestCase +from typing import Any, Dict, Optional, Sequence import pytest from common.aws import AWSInstance from monkey_island.cc.services import AWSService -from monkey_island.cc.services.aws_service import filter_instance_data_from_aws_response +EXPECTED_INSTANCE_1 = { + "instance_id": "1", + "name": "comp1", + "os": "linux", + "ip_address": "192.168.1.1", +} +EXPECTED_INSTANCE_2 = { + "instance_id": "2", + "name": "comp2", + "os": "linux", + "ip_address": "192.168.1.2", +} -class TestAwsService(TestCase): - def test_filter_instance_data_from_aws_response(self): - json_response_full = """ - { - "InstanceInformationList": [ - { - "ActivationId": "string", - "AgentVersion": "string", - "AssociationOverview": { - "DetailedStatus": "string", - "InstanceAssociationStatusAggregatedCount": { - "string" : 6 - } - }, - "AssociationStatus": "string", - "ComputerName": "string", - "IamRole": "string", - "InstanceId": "string", - "IPAddress": "string", - "IsLatestVersion": "True", - "LastAssociationExecutionDate": 6, - "LastPingDateTime": 6, - "LastSuccessfulAssociationExecutionDate": 6, - "Name": "string", - "PingStatus": "string", - "PlatformName": "string", - "PlatformType": "string", - "PlatformVersion": "string", - "RegistrationDate": 6, - "ResourceType": "string" - } - ], - "NextToken": "string" - } - """ - - json_response_empty = """ - { - "InstanceInformationList": [], - "NextToken": "string" - } - """ - - self.assertEqual( - filter_instance_data_from_aws_response(json.loads(json_response_empty)), [] - ) - self.assertEqual( - filter_instance_data_from_aws_response(json.loads(json_response_full)), - [{"instance_id": "string", "ip_address": "string", "name": "string", "os": "string"}], - ) +EMPTY_INSTANCE_INFO_RESPONSE = [] +FULL_INSTANCE_INFO_RESPONSE = [ + { + "ActivationId": "string", + "AgentVersion": "string", + "AssociationOverview": { + "DetailedStatus": "string", + "InstanceAssociationStatusAggregatedCount": {"string": 6}, + }, + "AssociationStatus": "string", + "ComputerName": EXPECTED_INSTANCE_1["name"], + "IamRole": "string", + "InstanceId": EXPECTED_INSTANCE_1["instance_id"], + "IPAddress": EXPECTED_INSTANCE_1["ip_address"], + "IsLatestVersion": "True", + "LastAssociationExecutionDate": 6, + "LastPingDateTime": 6, + "LastSuccessfulAssociationExecutionDate": 6, + "Name": "string", + "PingStatus": "string", + "PlatformName": "string", + "PlatformType": EXPECTED_INSTANCE_1["os"], + "PlatformVersion": "string", + "RegistrationDate": 6, + "ResourceType": "string", + }, + { + "ActivationId": "string", + "AgentVersion": "string", + "AssociationOverview": { + "DetailedStatus": "string", + "InstanceAssociationStatusAggregatedCount": {"string": 6}, + }, + "AssociationStatus": "string", + "ComputerName": EXPECTED_INSTANCE_2["name"], + "IamRole": "string", + "InstanceId": EXPECTED_INSTANCE_2["instance_id"], + "IPAddress": EXPECTED_INSTANCE_2["ip_address"], + "IsLatestVersion": "True", + "LastAssociationExecutionDate": 6, + "LastPingDateTime": 6, + "LastSuccessfulAssociationExecutionDate": 6, + "Name": "string", + "PingStatus": "string", + "PlatformName": "string", + "PlatformType": EXPECTED_INSTANCE_2["os"], + "PlatformVersion": "string", + "RegistrationDate": 6, + "ResourceType": "string", + }, +] class StubAWSInstance(AWSInstance): @@ -95,8 +105,12 @@ ACCOUNT_ID = "3" @pytest.fixture -def aws_service(): - aws_instance = StubAWSInstance(INSTANCE_ID, REGION, ACCOUNT_ID) +def aws_instance(): + return StubAWSInstance(INSTANCE_ID, REGION, ACCOUNT_ID) + + +@pytest.fixture +def aws_service(aws_instance): return AWSService(aws_instance) @@ -110,3 +124,27 @@ def test_region(aws_service): def test_account_id(aws_service): assert aws_service.island_aws_instance.account_id == ACCOUNT_ID + + +class MockAWSService(AWSService): + def __init__(self, aws_instance: AWSInstance, instance_info_response: Sequence[Dict[str, Any]]): + super().__init__(aws_instance) + self._instance_info_response = instance_info_response + + def _get_raw_managed_instances(self): + return self._instance_info_response + + +def test_get_managed_instances__empty(aws_instance): + aws_service = MockAWSService(aws_instance, EMPTY_INSTANCE_INFO_RESPONSE) + instances = aws_service.get_managed_instances() + assert len(instances) == 0 + + +def test_get_managed_instances(aws_instance): + aws_service = MockAWSService(aws_instance, FULL_INSTANCE_INFO_RESPONSE) + instances = aws_service.get_managed_instances() + + assert len(instances) == 2 + assert instances[0] == EXPECTED_INSTANCE_1 + assert instances[1] == EXPECTED_INSTANCE_2 From a0660f12e9ca7c2d7380dbeec32a30b5b6b02ec2 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 10:16:46 -0400 Subject: [PATCH 5/7] Island: Rename _filter_instance_info_from_aws_response --- monkey/monkey_island/cc/services/aws_service.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/services/aws_service.py b/monkey/monkey_island/cc/services/aws_service.py index d589fad96..6db1e7baf 100644 --- a/monkey/monkey_island/cc/services/aws_service.py +++ b/monkey/monkey_island/cc/services/aws_service.py @@ -28,7 +28,7 @@ class AWSService: def get_managed_instances(self) -> Sequence[Dict[str, str]]: raw_managed_instances_info = self._get_raw_managed_instances() - return _filter_instance_info_from_aws_response(raw_managed_instances_info) + return _filter_relevant_instance_info(raw_managed_instances_info) def _get_raw_managed_instances(self) -> Sequence[Dict[str, Any]]: """ @@ -57,7 +57,7 @@ class AWSService: pass -def _filter_instance_info_from_aws_response(raw_managed_instances_info: Sequence[Dict[str, Any]]): +def _filter_relevant_instance_info(raw_managed_instances_info: Sequence[Dict[str, Any]]): return [ { "instance_id": managed_instance[INSTANCE_ID_KEY], From 2da6e023e148e230262c1728e44c9220861fa3f7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 10:25:13 -0400 Subject: [PATCH 6/7] Island: Construct and register AWSService in the composition root --- monkey/monkey_island/cc/services/initialize.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index bcc11676d..8e0540329 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -1,9 +1,8 @@ from pathlib import Path -from threading import Thread from common import DIContainer from common.aws import AWSInstance -from monkey_island.cc.services import DirectoryFileStorageService, IFileStorageService, aws_service +from monkey_island.cc.services import AWSService, DirectoryFileStorageService, IFileStorageService from monkey_island.cc.services.post_breach_files import PostBreachFilesService from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService @@ -12,13 +11,12 @@ from . import AuthenticationService, JsonFileUserDatastore def initialize_services(data_dir: Path) -> DIContainer: container = DIContainer() + container.register_instance(AWSInstance, AWSInstance()) + container.register_instance( IFileStorageService, DirectoryFileStorageService(data_dir / "custom_pbas") ) - container.register_instance(AWSInstance, AWSInstance()) - - # Takes a while so it's best to start it in the background - Thread(target=aws_service.initialize, name="AwsService initialization", daemon=True).start() + container.register_instance(AWSService, container.resolve(AWSService)) # This is temporary until we get DI all worked out. PostBreachFilesService.initialize(container.resolve(IFileStorageService)) From cfbe1e565602313f8ab99189df7efcaf6491e284 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 12:45:40 -0400 Subject: [PATCH 7/7] Island: Add docstrings to AWSService --- .../monkey_island/cc/services/aws_service.py | 36 ++++++++++++++++--- 1 file changed, 32 insertions(+), 4 deletions(-) diff --git a/monkey/monkey_island/cc/services/aws_service.py b/monkey/monkey_island/cc/services/aws_service.py index 6db1e7baf..1a8dec455 100644 --- a/monkey/monkey_island/cc/services/aws_service.py +++ b/monkey/monkey_island/cc/services/aws_service.py @@ -1,5 +1,5 @@ import logging -from typing import Any, Dict, Iterable, Sequence +from typing import Any, Iterable, Mapping, Sequence import boto3 import botocore @@ -17,20 +17,37 @@ logger = logging.getLogger(__name__) class AWSService: def __init__(self, aws_instance: AWSInstance): + """ + :param aws_instance: An AWSInstance object representing the AWS instance that the Island is + running on + """ self._aws_instance = aws_instance def island_is_running_on_aws(self) -> bool: + """ + :return: True if the island is running on an AWS instance. False otherwise. + :rtype: bool + """ return self._aws_instance.is_instance @property def island_aws_instance(self) -> AWSInstance: + """ + :return: an AWSInstance object representing the AWS instance that the Island is running on. + :rtype: AWSInstance + """ return self._aws_instance - def get_managed_instances(self) -> Sequence[Dict[str, str]]: + def get_managed_instances(self) -> Sequence[Mapping[str, str]]: + """ + :return: A sequence of mappings, where each Mapping represents a managed AWS instance that + is accessible from the Island. + :rtype: Sequence[Mapping[str, str]] + """ raw_managed_instances_info = self._get_raw_managed_instances() return _filter_relevant_instance_info(raw_managed_instances_info) - def _get_raw_managed_instances(self) -> Sequence[Dict[str, Any]]: + def _get_raw_managed_instances(self) -> Sequence[Mapping[str, Any]]: """ Get the information for all instances with the relevant roles. @@ -57,7 +74,18 @@ class AWSService: pass -def _filter_relevant_instance_info(raw_managed_instances_info: Sequence[Dict[str, Any]]): +def _filter_relevant_instance_info(raw_managed_instances_info: Sequence[Mapping[str, Any]]): + """ + Consume raw instance data from the AWS API and return only those fields that are relevant for + Infection Monkey. + + :param raw_managed_instances_info: The output of + DescribeInstanceInformation["InstanceInformation"] from the + AWS API + :return: A sequence of mappings, where each Mapping represents a managed AWS instance that + is accessible from the Island. + :rtype: Sequence[Mapping[str, str]] + """ return [ { "instance_id": managed_instance[INSTANCE_ID_KEY],