forked from p15670423/monkey
Merge pull request #1639 from guardicore/1597-implement-stubs-in-automated-master
1597 implement stubs in automated master
This commit is contained in:
commit
84c6f6ee01
|
@ -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)
|
||||
|
|
|
@ -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/<string:guid>/<string:config_format>" route.
|
||||
# "/api/monkey/<string:guid>/<string:config_format>" route. Also considering not
|
||||
# flattening the config in the first place.
|
||||
if config_format == "legacy":
|
||||
ConfigService.decrypt_flat_config(monkey_json["config"])
|
||||
else:
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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": {
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue