Merge pull request #1639 from guardicore/1597-implement-stubs-in-automated-master

1597 implement stubs in automated master
This commit is contained in:
Mike Salvatore 2021-12-06 19:41:30 -05:00 committed by GitHub
commit 84c6f6ee01
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 138 additions and 44 deletions

View File

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

View File

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

View File

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

View File

@ -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": {

View File

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