Added generic collector processing functions, a dispatcher (name to function) with unit tests, and moved AWS to collector from regular sysinfo

This commit is contained in:
Shay Nehmad 2020-01-19 21:36:01 +02:00
parent 9583956683
commit 3496a78f6c
10 changed files with 165 additions and 29 deletions

View File

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

View File

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

View File

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

View File

@ -1,4 +1,4 @@
{
"server_config": "standard",
"server_config": "testing",
"deployment": "develop"
}

View File

@ -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."
},

View File

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

View File

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

View File

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

View File

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

View File

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