diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 1868aee1f..9c36dc17d 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -1,12 +1,13 @@ import logging import threading import time -from typing import Dict, List +from typing import Any, Callable, Dict, List, Tuple from infection_monkey.i_control_channel import IControlChannel from infection_monkey.i_master import IMaster from infection_monkey.i_puppet import IPuppet from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger +from infection_monkey.telemetry.post_breach_telem import PostBreachTelem from infection_monkey.telemetry.system_info_telem import SystemInfoTelem from infection_monkey.utils.timer import Timer @@ -29,8 +30,8 @@ class AutomatedMaster(IMaster): self._control_channel = control_channel self._stop = threading.Event() - self._master_thread = threading.Thread(target=self._run_master_thread, daemon=True) - self._simulation_thread = threading.Thread(target=self._run_simulation, daemon=True) + self._master_thread = _create_daemon_thread(target=self._run_master_thread) + self._simulation_thread = _create_daemon_thread(target=self._run_simulation) def start(self): logger.info("Starting automated breach and attack simulation") @@ -86,13 +87,17 @@ class AutomatedMaster(IMaster): def _run_simulation(self): config = self._control_channel.get_config() - system_info_collector_thread = threading.Thread( - target=self._collect_system_info, - args=(config["system_info_collector_classes"],), - daemon=True, + system_info_collector_thread = _create_daemon_thread( + target=self._run_plugins, + args=( + config["system_info_collector_classes"], + "system info collector", + self._collect_system_info, + ), ) - pba_thread = threading.Thread( - target=self._run_pbas, args=(config["post_breach_actions"],), daemon=True + pba_thread = _create_daemon_thread( + target=self._run_plugins, + args=(config["post_breach_actions"].items(), "post-breach action", self._run_pba), ) system_info_collector_thread.start() @@ -105,14 +110,13 @@ class AutomatedMaster(IMaster): system_info_collector_thread.join() if self._can_propagate(): - propagation_thread = threading.Thread( - target=self._propagate, args=(config,), daemon=True - ) + propagation_thread = _create_daemon_thread(target=self._propagate, args=(config,)) propagation_thread.start() propagation_thread.join() - payload_thread = threading.Thread( - target=self._run_payloads, args=(config["payloads"],), daemon=True + payload_thread = _create_daemon_thread( + target=self._run_plugins, + args=(config["payloads"].items(), "payload", self._run_payload), ) payload_thread.start() payload_thread.join() @@ -127,26 +131,19 @@ class AutomatedMaster(IMaster): if self._stop.is_set(): break - def _collect_system_info(self, enabled_collectors: List[str]): - logger.info("Running system info collectors") + def _collect_system_info(self, collector: str): + system_info_telemetry = {} + system_info_telemetry[collector] = self._puppet.run_sys_info_collector(collector) + self._telemetry_messenger.send_telemetry( + SystemInfoTelem({"collectors": system_info_telemetry}) + ) - for collector in enabled_collectors: - if self._stop.is_set(): - logger.debug("Received a stop signal, skipping remaining system info collectors") - break + def _run_pba(self, pba: Tuple[str, Dict]): + name = pba[0] + options = pba[1] - logger.info(f"Running system info collector: {collector}") - - system_info_telemetry = {} - system_info_telemetry[collector] = self._puppet.run_sys_info_collector(collector) - self._telemetry_messenger.send_telemetry( - SystemInfoTelem({"collectors": system_info_telemetry}) - ) - - logger.info("Finished running system info collectors") - - def _run_pbas(self, enabled_pbas: List[str]): - pass + command, result = self._puppet.run_pba(name, options) + self._telemetry_messenger.send_telemetry(PostBreachTelem(name, command, result)) def _can_propagate(self): return True @@ -154,8 +151,28 @@ class AutomatedMaster(IMaster): def _propagate(self, config: Dict): pass - def _run_payloads(self, enabled_payloads: Dict[str, Dict]): - pass + def _run_payload(self, payload: Tuple[str, Dict]): + name = payload[0] + options = payload[1] + + self._puppet.run_payload(name, options, self._stop) + + def _run_plugins(self, plugin: List[Any], plugin_type: str, callback: Callable[[Any], None]): + logger.info(f"Running {plugin_type}s") + logger.debug(f"Found {len(plugin)} {plugin_type}(s) to run") + + for p in plugin: + if self._stop.is_set(): + logger.debug(f"Received a stop signal, skipping remaining {plugin_type}s") + return + + callback(p) + + logger.info(f"Finished running {plugin_type}s") def cleanup(self): pass + + +def _create_daemon_thread(target: Callable[[Any], None], args: Tuple[Any] = ()): + return threading.Thread(target=target, args=args, daemon=True) diff --git a/monkey/monkey_island/cc/resources/monkey.py b/monkey/monkey_island/cc/resources/monkey.py index 4e02fc258..3853b58ed 100644 --- a/monkey/monkey_island/cc/resources/monkey.py +++ b/monkey/monkey_island/cc/resources/monkey.py @@ -27,7 +27,8 @@ class Monkey(flask_restful.Resource): if guid: monkey_json = mongo.db.monkey.find_one_or_404({"guid": guid}) # TODO: When the "legacy" format is no longer needed, update this logic and remove the - # "/api/monkey//" route. + # "/api/monkey//" route. Also considering not + # flattening the config in the first place. if config_format == "legacy": ConfigService.decrypt_flat_config(monkey_json["config"]) else: diff --git a/monkey/monkey_island/cc/services/config.py b/monkey/monkey_island/cc/services/config.py index 4e5290a19..97bbd4c82 100644 --- a/monkey/monkey_island/cc/services/config.py +++ b/monkey/monkey_island/cc/services/config.py @@ -430,6 +430,8 @@ class ConfigService: @staticmethod def format_flat_config_for_agent(config: Dict): ConfigService._remove_credentials_from_flat_config(config) + ConfigService._format_payloads_from_flat_config(config) + ConfigService._format_pbas_from_flat_config(config) @staticmethod def _remove_credentials_from_flat_config(config: Dict): @@ -443,3 +445,33 @@ class ConfigService: for field in fields_to_remove: config.pop(field, None) + + @staticmethod + def _format_payloads_from_flat_config(config: Dict): + config.setdefault("payloads", {})["ransomware"] = config["ransomware"] + config.pop("ransomware", None) + + @staticmethod + def _format_pbas_from_flat_config(config: Dict): + flat_linux_command_field = "custom_PBA_linux_cmd" + flat_linux_filename_field = "PBA_linux_filename" + flat_windows_command_field = "custom_PBA_windows_cmd" + flat_windows_filename_field = "PBA_windows_filename" + + formatted_pbas_config = {} + for pba in config.get("post_breach_actions", []): + formatted_pbas_config[pba] = {} + + formatted_pbas_config["Custom"] = { + "linux_command": config.get(flat_linux_command_field, ""), + "linux_filename": config.get(flat_linux_filename_field, ""), + "windows_command": config.get(flat_windows_command_field, ""), + "windows_filename": config.get(flat_windows_filename_field, ""), + } + + config["post_breach_actions"] = formatted_pbas_config + + config.pop(flat_linux_command_field, None) + config.pop(flat_linux_filename_field, None) + config.pop(flat_windows_command_field, None) + config.pop(flat_windows_filename_field, None) diff --git a/monkey/tests/data_for_tests/monkey_configs/flat_config.json b/monkey/tests/data_for_tests/monkey_configs/flat_config.json index 82cc895a1..b82ab6309 100644 --- a/monkey/tests/data_for_tests/monkey_configs/flat_config.json +++ b/monkey/tests/data_for_tests/monkey_configs/flat_config.json @@ -7,8 +7,8 @@ 7001, 9200 ], - "PBA_linux_filename": "", - "PBA_windows_filename": "", + "PBA_linux_filename": "test.sh", + "PBA_windows_filename": "test.ps1", "alive": true, "aws_access_key_id": "", "aws_secret_access_key": "", @@ -18,8 +18,8 @@ "10.197.94.72:5000" ], "current_server": "10.197.94.72:5000", - "custom_PBA_linux_cmd": "", - "custom_PBA_windows_cmd": "", + "custom_PBA_linux_cmd": "bash test.sh", + "custom_PBA_windows_cmd": "powershell test.ps1", "depth": 2, "dropper_date_reference_path_linux": "/bin/sh", "dropper_date_reference_path_windows": "%windir%\\system32\\kernel32.dll", @@ -82,9 +82,6 @@ "post_breach_actions": [ "CommunicateAsBackdoorUser", "ModifyShellStartupFiles", - "HiddenFiles", - "TrapCommand", - "ChangeSetuidSetgid", "ScheduleJobs", "Timestomping", "AccountDiscovery" @@ -93,8 +90,8 @@ "encryption": { "enabled": true, "directories": { - "linux_target_dir": "", - "windows_target_dir": "" + "linux_target_dir": "/tmp/ransomware-target", + "windows_target_dir": "C:\\windows\\temp\\ransomware-target" } }, "other_behaviors": { diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_config.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_config.py index 30e56e05e..be6bded05 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_config.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_config.py @@ -33,3 +33,50 @@ def test_format_config_for_agent__credentials_removed(flat_monkey_config): assert "exploit_password_list" not in flat_monkey_config assert "exploit_ssh_keys" not in flat_monkey_config assert "exploit_user_list" not in flat_monkey_config + + +def test_format_config_for_agent__ransomware_payload(flat_monkey_config): + expected_ransomware_config = { + "ransomware": { + "encryption": { + "enabled": True, + "directories": { + "linux_target_dir": "/tmp/ransomware-target", + "windows_target_dir": "C:\\windows\\temp\\ransomware-target", + }, + }, + "other_behaviors": {"readme": True}, + } + } + + ConfigService.format_flat_config_for_agent(flat_monkey_config) + + assert "payloads" in flat_monkey_config + assert flat_monkey_config["payloads"] == expected_ransomware_config + + assert "ransomware" not in flat_monkey_config + + +def test_format_config_for_agent__pbas(flat_monkey_config): + expected_pbas_config = { + "CommunicateAsBackdoorUser": {}, + "ModifyShellStartupFiles": {}, + "ScheduleJobs": {}, + "Timestomping": {}, + "AccountDiscovery": {}, + "Custom": { + "linux_command": "bash test.sh", + "windows_command": "powershell test.ps1", + "linux_filename": "test.sh", + "windows_filename": "test.ps1", + }, + } + ConfigService.format_flat_config_for_agent(flat_monkey_config) + + assert "post_breach_actions" in flat_monkey_config + assert flat_monkey_config["post_breach_actions"] == expected_pbas_config + + assert "custom_PBA_linux_cmd" not in flat_monkey_config + assert "PBA_linux_filename" not in flat_monkey_config + assert "custom_PBA_windows_cmd" not in flat_monkey_config + assert "PBA_windows_filename" not in flat_monkey_config