From 3496a78f6c0fdf2b0ee4128da2c14c9a77b72d6e Mon Sep 17 00:00:00 2001 From: Shay Nehmad Date: Sun, 19 Jan 2020 21:36:01 +0200 Subject: [PATCH] Added generic collector processing functions, a dispatcher (name to function) with unit tests, and moved AWS to collector from regular sysinfo --- monkey/infection_monkey/config.py | 2 +- .../infection_monkey/system_info/__init__.py | 10 --- .../system_info/collectors/aws_collector.py | 30 +++++++++ monkey/monkey_island/cc/server_config.json | 2 +- .../cc/services/config_schema.py | 11 +++- .../telemetry/processing/system_info.py | 11 +--- .../processing/system_info_collectors/aws.py | 15 +++++ .../system_info_collectors/environment.py | 12 ++-- .../system_info_telemetry_dispatcher.py | 36 ++++++++++ .../test_system_info_telemetry_dispatcher.py | 65 +++++++++++++++++++ 10 files changed, 165 insertions(+), 29 deletions(-) create mode 100644 monkey/infection_monkey/system_info/collectors/aws_collector.py create mode 100644 monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/aws.py create mode 100644 monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py create mode 100644 monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py diff --git a/monkey/infection_monkey/config.py b/monkey/infection_monkey/config.py index e76ed8101..e1b1ece83 100644 --- a/monkey/infection_monkey/config.py +++ b/monkey/infection_monkey/config.py @@ -125,7 +125,7 @@ class Configuration(object): finger_classes = [] exploiter_classes = [] - system_info_collectors_classes = ["EnvironmentCollector"] + system_info_collectors_classes = ["EnvironmentCollector", "AwsCollector"] # how many victims to look for in a single scan iteration victims_max_find = 100 diff --git a/monkey/infection_monkey/system_info/__init__.py b/monkey/infection_monkey/system_info/__init__.py index ccec45cde..fc4aa1caf 100644 --- a/monkey/infection_monkey/system_info/__init__.py +++ b/monkey/infection_monkey/system_info/__init__.py @@ -6,7 +6,6 @@ import psutil from enum import IntEnum from infection_monkey.network.info import get_host_subnets -from infection_monkey.system_info.aws_collector import AwsCollector from infection_monkey.system_info.azure_cred_collector import AzureCollector from infection_monkey.system_info.netstat_collector import NetstatCollector from system_info.system_info_collectors_handler import SystemInfoCollectorsHandler @@ -67,7 +66,6 @@ class InfoCollector(object): self.get_process_list() self.get_network_info() self.get_azure_info() - self.get_aws_info() # Collect all plugins SystemInfoCollectorsHandler().execute_all_configured() @@ -155,11 +153,3 @@ class InfoCollector(object): except Exception: # If we failed to collect azure info, no reason to fail all the collection. Log and continue. LOG.error("Failed collecting Azure info.", exc_info=True) - - def get_aws_info(self): - # noinspection PyBroadException - try: - self.info['aws'] = AwsCollector().get_aws_info() - except Exception: - # If we failed to collect aws info, no reason to fail all the collection. Log and continue. - LOG.error("Failed collecting AWS info.", exc_info=True) diff --git a/monkey/infection_monkey/system_info/collectors/aws_collector.py b/monkey/infection_monkey/system_info/collectors/aws_collector.py new file mode 100644 index 000000000..71f9e58c1 --- /dev/null +++ b/monkey/infection_monkey/system_info/collectors/aws_collector.py @@ -0,0 +1,30 @@ +import logging + +from common.cloud.aws.aws_instance import AwsInstance +from infection_monkey.system_info.system_info_collector import SystemInfoCollector + + +logger = logging.getLogger(__name__) + + +class AwsCollector(SystemInfoCollector): + """ + Extract info from AWS machines. + """ + def __init__(self): + super(AwsCollector, self).__init__(name="AwsCollector") + + def collect(self) -> dict: + logger.info("Collecting AWS info") + aws = AwsInstance() + info = {} + if aws.is_instance(): + logger.info("Machine is an AWS instance") + info = \ + { + 'instance_id': aws.get_instance_id() + } + else: + logger.info("Machine is NOT an AWS instance") + + return info diff --git a/monkey/monkey_island/cc/server_config.json b/monkey/monkey_island/cc/server_config.json index 420f1b303..7bf106194 100644 --- a/monkey/monkey_island/cc/server_config.json +++ b/monkey/monkey_island/cc/server_config.json @@ -1,4 +1,4 @@ { - "server_config": "standard", + "server_config": "testing", "deployment": "develop" } diff --git a/monkey/monkey_island/cc/services/config_schema.py b/monkey/monkey_island/cc/services/config_schema.py index 5de57e26b..d5e015866 100644 --- a/monkey/monkey_island/cc/services/config_schema.py +++ b/monkey/monkey_island/cc/services/config_schema.py @@ -111,6 +111,14 @@ SCHEMA = { "title": "Which Environment this machine is on (on prem/cloud)", "attack_techniques": [] }, + { + "type": "string", + "enum": [ + "AwsCollector" + ], + "title": "If on AWS, collect more information about the instance", + "attack_techniques": [] + }, ], }, "post_breach_acts": { @@ -455,7 +463,8 @@ SCHEMA = { "$ref": "#/definitions/system_info_collectors_classes" }, "default": [ - "EnvironmentCollector" + "EnvironmentCollector", + "AwsCollector" ], "description": "Determines which system information collectors will collect information." }, diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info.py index 04ab27d95..915fa7a25 100644 --- a/monkey/monkey_island/cc/services/telemetry/processing/system_info.py +++ b/monkey/monkey_island/cc/services/telemetry/processing/system_info.py @@ -5,6 +5,7 @@ from monkey_island.cc.models import Monkey from monkey_island.cc.services import mimikatz_utils from monkey_island.cc.services.node import NodeService from monkey_island.cc.services.config import ConfigService +from monkey_island.cc.services.telemetry.processing.system_info_collectors.aws import process_aws_telemetry from monkey_island.cc.services.telemetry.processing.system_info_collectors.environment import process_environment_telemetry from monkey_island.cc.services.telemetry.zero_trust_tests.antivirus_existence import test_antivirus_existence from monkey_island.cc.services.wmi_handler import WMIHandler @@ -18,7 +19,7 @@ def process_system_info_telemetry(telemetry_json): process_ssh_info, process_credential_info, process_mimikatz_and_wmi_info, - process_aws_data, + process_aws_telemetry, update_db_with_new_hostname, test_antivirus_existence, process_environment_telemetry @@ -116,13 +117,5 @@ def process_mimikatz_and_wmi_info(telemetry_json): wmi_handler.process_and_handle_wmi_info() -def process_aws_data(telemetry_json): - if 'aws' in telemetry_json['data']: - if 'instance_id' in telemetry_json['data']['aws']: - monkey_id = NodeService.get_monkey_by_guid(telemetry_json['monkey_guid']).get('_id') - mongo.db.monkey.update_one({'_id': monkey_id}, - {'$set': {'aws_instance_id': telemetry_json['data']['aws']['instance_id']}}) - - def update_db_with_new_hostname(telemetry_json): Monkey.get_single_monkey_by_guid(telemetry_json['monkey_guid']).set_hostname(telemetry_json['data']['hostname']) diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/aws.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/aws.py new file mode 100644 index 000000000..2b4d8085e --- /dev/null +++ b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/aws.py @@ -0,0 +1,15 @@ +import logging + +from monkey_island.cc.models.monkey import Monkey + +logger = logging.getLogger(__name__) + + +def process_aws_telemetry(collector_results, monkey_guid): + relevant_monkey = Monkey.get_single_monkey_by_guid(monkey_guid) + + if "instance_id" in collector_results: + instance_id = collector_results["instance_id"] + relevant_monkey.aws_instance_id = instance_id + relevant_monkey.save() + logger.debug("Updated Monkey {} with aws instance id {}".format(str(relevant_monkey), instance_id)) diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/environment.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/environment.py index d66019b39..9ddee70ce 100644 --- a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/environment.py +++ b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/environment.py @@ -5,10 +5,8 @@ from monkey_island.cc.models.monkey import Monkey logger = logging.getLogger(__name__) -def process_environment_telemetry(telemetry_json): - if "EnvironmentCollector" in telemetry_json["data"]["collectors"]: - env = telemetry_json["data"]["collectors"]["EnvironmentCollector"]["environment"] - relevant_monkey = Monkey.get_single_monkey_by_guid(telemetry_json['monkey_guid']) - relevant_monkey.environment = env - relevant_monkey.save() - logger.debug("Updated Monkey {} with env {}".format(str(relevant_monkey), env)) +def process_environment_telemetry(collector_results, monkey_guid): + relevant_monkey = Monkey.get_single_monkey_by_guid(monkey_guid) + relevant_monkey.environment = collector_results + relevant_monkey.save() + logger.debug("Updated Monkey {} with env {}".format(str(relevant_monkey), collector_results)) diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py new file mode 100644 index 000000000..661034efb --- /dev/null +++ b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py @@ -0,0 +1,36 @@ +import logging + +from monkey_island.cc.services.telemetry.processing.system_info_collectors.aws import process_aws_telemetry +from monkey_island.cc.services.telemetry.processing.system_info_collectors.environment import process_environment_telemetry + +logger = logging.getLogger(__name__) + +SYSTEM_INFO_COLLECTOR_TO_TELEMETRY_PROCESSOR = { + "AwsCollector": process_aws_telemetry, + "EnvironmentCollector": process_environment_telemetry, +} + + +class SystemInfoTelemetryDispatcher(object): + def __init__(self, collector_to_parsing_function=None): + if collector_to_parsing_function is None: + collector_to_parsing_function = SYSTEM_INFO_COLLECTOR_TO_TELEMETRY_PROCESSOR + self.collector_to_parsing_function = collector_to_parsing_function + + def dispatch_to_relevant_collector(self, telemetry_json): + if "collectors" in telemetry_json["data"]: + self.send_each_result_to_relevant_processor(telemetry_json) + + def send_each_result_to_relevant_processor(self, telemetry_json): + relevant_monkey_guid = telemetry_json['monkey_guid'] + for collector_name, collector_results in telemetry_json["data"]["collectors"].items(): + if collector_name in self.collector_to_parsing_function: + # noinspection PyBroadException + try: + self.collector_to_parsing_function[collector_name](collector_results, relevant_monkey_guid) + except Exception as e: + logger.error( + "Error {} while processing {} system info telemetry".format(str(e), collector_name), + exc_info=True) + else: + logger.warning("Unknown system info collector name: {}".format(collector_name)) diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py new file mode 100644 index 000000000..db36cd5bb --- /dev/null +++ b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py @@ -0,0 +1,65 @@ +from importlib import reload +from unittest import mock +from unittest.mock import MagicMock + +import uuid + +from monkey_island.cc.models import Monkey +from monkey_island.cc.services.telemetry.processing.system_info_collectors.system_info_telemetry_dispatcher import \ + SystemInfoTelemetryDispatcher +from monkey_island.cc.testing.IslandTestCase import IslandTestCase +from monkey_island.cc.services.telemetry.processing.system_info_collectors.system_info_telemetry_dispatcher import \ + process_aws_telemetry + +TEST_SYS_INFO_TO_PROCESSING = { + "AwsCollector": process_aws_telemetry, +} + + +def do_nothing(x, y): + pass + + +class SystemInfoTelemetryDispatcherTest(IslandTestCase): + def test_dispatch_to_relevant_collector_bad_inputs(self): + self.fail_if_not_testing_env() + + dispatcher = SystemInfoTelemetryDispatcher(TEST_SYS_INFO_TO_PROCESSING) + + # Bad format telem JSONs - throws + bad_empty_telem_json = {} + self.assertRaises(KeyError, dispatcher.dispatch_to_relevant_collector, bad_empty_telem_json) + bad_no_data_telem_json = {"monkey_guid": "bla"} + self.assertRaises(KeyError, dispatcher.dispatch_to_relevant_collector, bad_no_data_telem_json) + bad_no_monkey_telem_json = {"data": {"collectors": {"AwsCollector": "Bla"}}} + self.assertRaises(KeyError, dispatcher.dispatch_to_relevant_collector, bad_no_monkey_telem_json) + + # Telem JSON with no collectors - nothing gets dispatched + good_telem_no_collectors = {"monkey_guid": "bla", "data": {"bla": "bla"}} + good_telem_empty_collectors = {"monkey_guid": "bla", "data": {"bla": "bla", "collectors": {}}} + + dispatcher.dispatch_to_relevant_collector(good_telem_no_collectors) + dispatcher.dispatch_to_relevant_collector(good_telem_empty_collectors) + + def test_dispatch_to_relevant_collector(self): + self.fail_if_not_testing_env() + self.clean_monkey_db() + + a_monkey = Monkey(guid=str(uuid.uuid4())) + a_monkey.save() + + dispatcher = SystemInfoTelemetryDispatcher() + + # JSON with results - make sure functions are called + instance_id = "i-0bd2c14bd4c7d703f" + telem_json = { + "data": { + "collectors": { + "AwsCollector": {"instance_id": instance_id}, + } + }, + "monkey_guid": a_monkey.guid + } + dispatcher.dispatch_to_relevant_collector(telem_json) + + self.assertEquals(Monkey.get_single_monkey_by_guid(a_monkey.guid).aws_instance_id, instance_id)