From 44055b32f9c9fb4981d3f92b02595e06a7d6a857 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Fri, 3 Dec 2021 09:17:33 -0500
Subject: [PATCH 1/7] Island: Reformat "payloads" in config before sending to
 agent

Allow the configuration to contain multiple payloads that can be run by
the agent.
---
 monkey/monkey_island/cc/services/config.py    |  6 +++++
 .../monkey_configs/flat_config.json           |  4 ++--
 .../monkey_island/cc/services/test_config.py  | 22 +++++++++++++++++++
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/monkey/monkey_island/cc/services/config.py b/monkey/monkey_island/cc/services/config.py
index 4e5290a19..80228c8e6 100644
--- a/monkey/monkey_island/cc/services/config.py
+++ b/monkey/monkey_island/cc/services/config.py
@@ -430,6 +430,7 @@ 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)
 
     @staticmethod
     def _remove_credentials_from_flat_config(config: Dict):
@@ -443,3 +444,8 @@ 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)
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..1f700d40f 100644
--- a/monkey/tests/data_for_tests/monkey_configs/flat_config.json
+++ b/monkey/tests/data_for_tests/monkey_configs/flat_config.json
@@ -93,8 +93,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..2f67c2f76 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,25 @@ 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

From 839157a8226cd2bded6c23685bf69964a4f084b8 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Fri, 3 Dec 2021 09:39:41 -0500
Subject: [PATCH 2/7] Agent: Implement AutomatedMaster._run_payloads()

---
 monkey/infection_monkey/master/automated_master.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py
index 1868aee1f..d50f242c1 100644
--- a/monkey/infection_monkey/master/automated_master.py
+++ b/monkey/infection_monkey/master/automated_master.py
@@ -155,7 +155,17 @@ class AutomatedMaster(IMaster):
         pass
 
     def _run_payloads(self, enabled_payloads: Dict[str, Dict]):
-        pass
+        logger.info("Running payloads")
+        logger.debug(f"Found {len(enabled_payloads.keys())} payload(s) to run")
+
+        for payload_name, options in enabled_payloads.items():
+            if self._stop.is_set():
+                logger.debug("Received a stop signal, skipping remaining system info collectors")
+                break
+
+            self._puppet.run_payload(payload_name, options, self._stop)
+
+        logger.info("Finished running payloads")
 
     def cleanup(self):
         pass

From 1b04844e5efa13fe46172365be4d34400d234f5d Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Fri, 3 Dec 2021 10:21:10 -0500
Subject: [PATCH 3/7] Agent: Deduplicate stop logic in AutomatedMaster

---
 .../master/automated_master.py                | 59 ++++++++++---------
 1 file changed, 30 insertions(+), 29 deletions(-)

diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py
index d50f242c1..42ab4d285 100644
--- a/monkey/infection_monkey/master/automated_master.py
+++ b/monkey/infection_monkey/master/automated_master.py
@@ -1,7 +1,7 @@
 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
@@ -87,8 +87,12 @@ class AutomatedMaster(IMaster):
         config = self._control_channel.get_config()
 
         system_info_collector_thread = threading.Thread(
-            target=self._collect_system_info,
-            args=(config["system_info_collector_classes"],),
+            target=self._run_plugins,
+            args=(
+                config["system_info_collector_classes"],
+                "system info collector",
+                self._collect_system_info,
+            ),
             daemon=True,
         )
         pba_thread = threading.Thread(
@@ -112,7 +116,9 @@ class AutomatedMaster(IMaster):
             propagation_thread.join()
 
         payload_thread = threading.Thread(
-            target=self._run_payloads, args=(config["payloads"],), daemon=True
+            target=self._run_plugins,
+            args=(config["payloads"].items(), "payload", self._run_payload),
+            daemon=True,
         )
         payload_thread.start()
         payload_thread.join()
@@ -127,23 +133,12 @@ class AutomatedMaster(IMaster):
             if self._stop.is_set():
                 break
 
-    def _collect_system_info(self, enabled_collectors: List[str]):
-        logger.info("Running system info collectors")
-
-        for collector in enabled_collectors:
-            if self._stop.is_set():
-                logger.debug("Received a stop signal, skipping remaining system info collectors")
-                break
-
-            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 _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})
+        )
 
     def _run_pbas(self, enabled_pbas: List[str]):
         pass
@@ -154,18 +149,24 @@ class AutomatedMaster(IMaster):
     def _propagate(self, config: Dict):
         pass
 
-    def _run_payloads(self, enabled_payloads: Dict[str, Dict]):
-        logger.info("Running payloads")
-        logger.debug(f"Found {len(enabled_payloads.keys())} payload(s) to run")
+    def _run_payload(self, payload: Tuple[str, Dict]):
+        name = payload[0]
+        options = payload[1]
 
-        for payload_name, options in enabled_payloads.items():
+        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("Received a stop signal, skipping remaining system info collectors")
-                break
+                logger.debug(f"Received a stop signal, skipping remaining {plugin_type}s")
+                return
 
-            self._puppet.run_payload(payload_name, options, self._stop)
+            callback(p)
 
-        logger.info("Finished running payloads")
+        logger.info(f"Finished running {plugin_type}s")
 
     def cleanup(self):
         pass

From fecb7342ade16e1b3734c3e30defd1d11efce0fc Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Fri, 3 Dec 2021 10:49:56 -0500
Subject: [PATCH 4/7] Island: Reformat "PBAs" in config before sending to agent

Allow options to be specified for each PBA and consolidate the custom
user PBA options under a "Custom" PBA.
---
 monkey/monkey_island/cc/services/config.py    | 26 +++++++++++++++++++
 .../monkey_configs/flat_config.json           | 11 +++-----
 .../monkey_island/cc/services/test_config.py  | 25 ++++++++++++++++++
 3 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/monkey/monkey_island/cc/services/config.py b/monkey/monkey_island/cc/services/config.py
index 80228c8e6..97bbd4c82 100644
--- a/monkey/monkey_island/cc/services/config.py
+++ b/monkey/monkey_island/cc/services/config.py
@@ -431,6 +431,7 @@ class ConfigService:
     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):
@@ -449,3 +450,28 @@ class ConfigService:
     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 1f700d40f..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"
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 2f67c2f76..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
@@ -55,3 +55,28 @@ def test_format_config_for_agent__ransomware_payload(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

From 261826fc787278f91fdc4f6a7133852f7e63693c Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Fri, 3 Dec 2021 11:05:31 -0500
Subject: [PATCH 5/7] Agent: Implement PBA thread in AutomatedMaster

---
 monkey/infection_monkey/master/automated_master.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py
index 42ab4d285..f0e17b8a2 100644
--- a/monkey/infection_monkey/master/automated_master.py
+++ b/monkey/infection_monkey/master/automated_master.py
@@ -7,6 +7,7 @@ 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
 
@@ -96,7 +97,9 @@ class AutomatedMaster(IMaster):
             daemon=True,
         )
         pba_thread = threading.Thread(
-            target=self._run_pbas, args=(config["post_breach_actions"],), daemon=True
+            target=self._run_plugins,
+            args=(config["post_breach_actions"].items(), "post-breach action", self._run_pba),
+            daemon=True,
         )
 
         system_info_collector_thread.start()
@@ -140,8 +143,12 @@ class AutomatedMaster(IMaster):
             SystemInfoTelem({"collectors": system_info_telemetry})
         )
 
-    def _run_pbas(self, enabled_pbas: List[str]):
-        pass
+    def _run_pba(self, pba: Tuple[str, Dict]):
+        name = pba[0]
+        options = pba[1]
+
+        command, result = self._puppet.run_pba(name, options)
+        self._telemetry_messenger.send_telemetry(PostBreachTelem(name, command, result))
 
     def _can_propagate(self):
         return True

From e8de38881c0b7749fa04feffbba1b873b2af4444 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Mon, 6 Dec 2021 19:13:53 -0500
Subject: [PATCH 6/7] Agent: Add _create_daemon_thread() utility function to
 AutomatedMaster

---
 .../master/automated_master.py                | 21 +++++++++----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py
index f0e17b8a2..9c36dc17d 100644
--- a/monkey/infection_monkey/master/automated_master.py
+++ b/monkey/infection_monkey/master/automated_master.py
@@ -30,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")
@@ -87,19 +87,17 @@ class AutomatedMaster(IMaster):
     def _run_simulation(self):
         config = self._control_channel.get_config()
 
-        system_info_collector_thread = threading.Thread(
+        system_info_collector_thread = _create_daemon_thread(
             target=self._run_plugins,
             args=(
                 config["system_info_collector_classes"],
                 "system info collector",
                 self._collect_system_info,
             ),
-            daemon=True,
         )
-        pba_thread = threading.Thread(
+        pba_thread = _create_daemon_thread(
             target=self._run_plugins,
             args=(config["post_breach_actions"].items(), "post-breach action", self._run_pba),
-            daemon=True,
         )
 
         system_info_collector_thread.start()
@@ -112,16 +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(
+        payload_thread = _create_daemon_thread(
             target=self._run_plugins,
             args=(config["payloads"].items(), "payload", self._run_payload),
-            daemon=True,
         )
         payload_thread.start()
         payload_thread.join()
@@ -177,3 +172,7 @@ class AutomatedMaster(IMaster):
 
     def cleanup(self):
         pass
+
+
+def _create_daemon_thread(target: Callable[[Any], None], args: Tuple[Any] = ()):
+    return threading.Thread(target=target, args=args, daemon=True)

From b15612c9aeb8b4c64abbabf78c717bc457b9a403 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Mon, 6 Dec 2021 19:31:50 -0500
Subject: [PATCH 7/7] Island: Add more detail to TODO in Monkey resource

---
 monkey/monkey_island/cc/resources/monkey.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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