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