Merge pull request #1818 from guardicore/1604-custom-pba

1604 custom pba

monkey/infection_monkey/config.py

@@ -89,15 +89,6 @@ class Configuration(object):
 
     keep_tunnel_open_time = 60
 
-    ###########################
-    # post breach actions
-    ###########################
-    post_breach_actions = []
-    custom_PBA_linux_cmd = ""
-    custom_PBA_windows_cmd = ""
-    PBA_linux_filename = None
-    PBA_windows_filename = None
-
     ###########################
     # testing configuration
     ###########################

monkey/infection_monkey/master/automated_master.py

@@ -1,7 +1,7 @@
 import logging
 import threading
 import time
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
+from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple
 
 from infection_monkey.credential_store import ICredentialsStore
 from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
@@ -16,6 +16,7 @@ from infection_monkey.utils.threading import create_daemon_thread, interruptible
 from infection_monkey.utils.timer import Timer
 
 from . import Exploiter, IPScanner, Propagator
+from .option_parsing import custom_pba_is_enabled
 
 CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC = 5
 CHECK_FOR_TERMINATE_INTERVAL_SEC = CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC / 5
@@ -154,9 +155,9 @@ class AutomatedMaster(IMaster):
             ),
         )
         pba_thread = create_daemon_thread(
-            target=self._run_plugins,
+            target=self._run_pbas,
             name="PBAThread",
-            args=(config["post_breach_actions"].items(), "post-breach action", self._run_pba),
+            args=(config["post_breach_actions"].items(), self._run_pba, config["custom_pbas"]),
         )
 
         credential_collector_thread.start()
@@ -196,10 +197,6 @@ class AutomatedMaster(IMaster):
         name = pba[0]
         options = pba[1]
 
-        # TEMPORARY; TO AVOID ERRORS SINCE THIS ISN'T IMPLEMENTED YET
-        if name == "Custom":
-            return
-
         for pba_data in self._puppet.run_pba(name, options):
             self._telemetry_messenger.send_telemetry(PostBreachTelem(pba_data))
 
@@ -212,6 +209,14 @@ class AutomatedMaster(IMaster):
 
         self._puppet.run_payload(name, options, self._stop)
 
+    def _run_pbas(
+        self, plugins: Iterable[Any], callback: Callable[[Any], None], custom_pba_options: Mapping
+    ):
+        self._run_plugins(plugins, "post-breach action", callback)
+
+        if custom_pba_is_enabled(custom_pba_options):
+            self._run_plugins([("CustomPBA", custom_pba_options)], "post-breach action", callback)
+
     def _run_plugins(
         self, plugins: Iterable[Any], plugin_type: str, callback: Callable[[Any], None]
     ):

monkey/infection_monkey/master/option_parsing.py

@@ -0,0 +1,13 @@
+from typing import Dict
+
+from infection_monkey.utils.environment import is_windows_os
+
+
+def custom_pba_is_enabled(pba_options: Dict) -> bool:
+    if not is_windows_os():
+        if pba_options["linux_command"]:
+            return True
+    else:
+        if pba_options["windows_command"]:
+            return True
+    return False

monkey/infection_monkey/monkey.py

@@ -51,6 +51,7 @@ from infection_monkey.post_breach.actions.schedule_jobs import ScheduleJobs
 from infection_monkey.post_breach.actions.timestomping import Timestomping
 from infection_monkey.post_breach.actions.use_signed_scripts import SignedScriptProxyExecution
 from infection_monkey.post_breach.actions.use_trap_command import TrapCommand
+from infection_monkey.post_breach.custom_pba import CustomPBA
 from infection_monkey.puppet.puppet import Puppet
 from infection_monkey.system_singleton import SystemSingleton
 from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
@@ -315,6 +316,9 @@ class InfectionMonkey:
             ClearCommandHistory(self._telemetry_messenger),
             PluginType.POST_BREACH_ACTION,
         )
+        puppet.load_plugin(
+            "CustomPBA", CustomPBA(self._telemetry_messenger), PluginType.POST_BREACH_ACTION
+        )
 
         puppet.load_plugin("ransomware", RansomwarePayload(), PluginType.PAYLOAD)
 

monkey/infection_monkey/post_breach/custom_pba/__init__.py

@@ -0,0 +1 @@
+from .custom_pba import CustomPBA

monkey/infection_monkey/post_breach/custom_pba/custom_pba.py

@@ -1,10 +1,11 @@
 import logging
 import os
+from typing import Dict, Iterable
 
 from common.common_consts.post_breach_consts import POST_BREACH_FILE_EXECUTION
 from common.utils.attack_utils import ScanStatus
-from infection_monkey.config import WormConfiguration
 from infection_monkey.control import ControlClient
+from infection_monkey.i_puppet import PostBreachData
 from infection_monkey.network.tools import get_interface_to_target
 from infection_monkey.post_breach.pba import PBA
 from infection_monkey.telemetry.attack.t1105_telem import T1105Telem
@@ -19,52 +20,50 @@ DIR_CHANGE_WINDOWS = "cd %s & "
 DIR_CHANGE_LINUX = "cd %s ; "
 
 
-class UsersPBA(PBA):
+class CustomPBA(PBA):
     """
     Defines user's configured post breach action.
     """
 
     def __init__(self, telemetry_messenger: ITelemetryMessenger):
-        super(UsersPBA, self).__init__(telemetry_messenger, POST_BREACH_FILE_EXECUTION)
+        super(CustomPBA, self).__init__(telemetry_messenger, POST_BREACH_FILE_EXECUTION)
         self.filename = ""
 
-        if not is_windows_os():
+    def run(self, options: Dict) -> Iterable[PostBreachData]:
-            # Add linux commands to PBA's
+        self._set_options(options)
-            if WormConfiguration.PBA_linux_filename:
+        return super().run(options)
-                self.filename = WormConfiguration.PBA_linux_filename
+
-                if WormConfiguration.custom_PBA_linux_cmd:
+    def _set_options(self, options: Dict):
-                    # Add change dir command, because user will try to access his file
+        # Required for attack telemetry
-                    self.command = (
+        self.current_server = options["current_server"]
-                        DIR_CHANGE_LINUX % get_monkey_dir_path()
+
-                    ) + WormConfiguration.custom_PBA_linux_cmd
+        if is_windows_os():
-            elif WormConfiguration.custom_PBA_linux_cmd:
-                self.command = WormConfiguration.custom_PBA_linux_cmd
-        else:
             # Add windows commands to PBA's
-            if WormConfiguration.PBA_windows_filename:
+            if options["windows_filename"]:
-                self.filename = WormConfiguration.PBA_windows_filename
+                self.filename = options["windows_filename"]
-                if WormConfiguration.custom_PBA_windows_cmd:
+                if options["windows_command"]:
                     # Add change dir command, because user will try to access his file
-                    self.command = (
+                    self.command = (DIR_CHANGE_WINDOWS % get_monkey_dir_path()) + options[
-                        DIR_CHANGE_WINDOWS % get_monkey_dir_path()
+                        "windows_command"
-                    ) + WormConfiguration.custom_PBA_windows_cmd
+                    ]
-            elif WormConfiguration.custom_PBA_windows_cmd:
+            elif options["windows_command"]:
-                self.command = WormConfiguration.custom_PBA_windows_cmd
+                self.command = options["windows_command"]
+        else:
+            # Add linux commands to PBA's
+            if options["linux_filename"]:
+                self.filename = options["linux_filename"]
+                if options["linux_command"]:
+                    # Add change dir command, because user will try to access his file
+                    self.command = (DIR_CHANGE_LINUX % get_monkey_dir_path()) + options[
+                        "linux_command"
+                    ]
+            elif options["linux_command"]:
+                self.command = options["linux_command"]
 
     def _execute_default(self):
         if self.filename:
-            UsersPBA.download_pba_file(get_monkey_dir_path(), self.filename)
+            self.download_pba_file(get_monkey_dir_path(), self.filename)
-        return super(UsersPBA, self)._execute_default()
+        return super(CustomPBA, self)._execute_default()
-
-    @staticmethod
-    def should_run(class_name):
-        if not is_windows_os():
-            if WormConfiguration.PBA_linux_filename or WormConfiguration.custom_PBA_linux_cmd:
-                return True
-        else:
-            if WormConfiguration.PBA_windows_filename or WormConfiguration.custom_PBA_windows_cmd:
-                return True
-        return False
 
     def download_pba_file(self, dst_dir, filename):
         """
@@ -84,11 +83,11 @@ class UsersPBA(PBA):
         if not status:
             status = ScanStatus.USED
 
-        self._telemetry_messenger.send_telemetry(
+        self.telemetry_messenger.send_telemetry(
             T1105Telem(
                 status,
-                WormConfiguration.current_server.split(":")[0],
+                self.current_server.split(":")[0],
-                get_interface_to_target(WormConfiguration.current_server.split(":")[0]),
+                get_interface_to_target(self.current_server.split(":")[0]),
                 filename,
             )
         )

monkey/monkey_island/cc/services/config.py

@@ -452,11 +452,13 @@ class ConfigService:
         for pba in config.get("post_breach_actions", []):
             formatted_pbas_config[pba] = {}
 
-        formatted_pbas_config["Custom"] = {
+        config["custom_pbas"] = {
             "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, ""),
+            # Current server is used for attack telemetry
+            "current_server": config.get("current_server"),
         }
 
         config["post_breach_actions"] = formatted_pbas_config

monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py

@@ -2,19 +2,20 @@ from unittest.mock import MagicMock
 
 import pytest
 
-from infection_monkey.post_breach.actions.users_custom_pba import UsersPBA
+from infection_monkey.post_breach.custom_pba.custom_pba import CustomPBA
 
 MONKEY_DIR_PATH = "/dir/to/monkey/"
 CUSTOM_LINUX_CMD = "command-for-linux"
 CUSTOM_LINUX_FILENAME = "filename-for-linux"
 CUSTOM_WINDOWS_CMD = "command-for-windows"
 CUSTOM_WINDOWS_FILENAME = "filename-for-windows"
+CUSTOM_SERVER = "10.10.10.10:5000"
 
 
-@pytest.fixture
+@pytest.fixture(autouse=True)
 def fake_monkey_dir_path(monkeypatch):
     monkeypatch.setattr(
-        "infection_monkey.post_breach.actions.users_custom_pba.get_monkey_dir_path",
+        "infection_monkey.post_breach.custom_pba.custom_pba.get_monkey_dir_path",
         lambda: MONKEY_DIR_PATH,
     )
 
@@ -22,7 +23,7 @@ def fake_monkey_dir_path(monkeypatch):
 @pytest.fixture
 def set_os_linux(monkeypatch):
     monkeypatch.setattr(
-        "infection_monkey.post_breach.actions.users_custom_pba.is_windows_os",
+        "infection_monkey.post_breach.custom_pba.custom_pba.is_windows_os",
         lambda: False,
     )
 
@@ -30,106 +31,92 @@ def set_os_linux(monkeypatch):
 @pytest.fixture
 def set_os_windows(monkeypatch):
     monkeypatch.setattr(
-        "infection_monkey.post_breach.actions.users_custom_pba.is_windows_os",
+        "infection_monkey.post_breach.custom_pba.custom_pba.is_windows_os",
         lambda: True,
     )
 
 
 @pytest.fixture
-def mock_UsersPBA_linux_custom_file_and_cmd(set_os_linux, fake_monkey_dir_path, monkeypatch):
+def fake_custom_pba_linux_options():
-    monkeypatch.setattr(
+    return {
-        "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd",
+        "linux_command": CUSTOM_LINUX_CMD,
-        CUSTOM_LINUX_CMD,
+        "linux_filename": CUSTOM_LINUX_FILENAME,
-    )
+        "windows_command": "",
-    monkeypatch.setattr(
+        "windows_filename": "",
-        "infection_monkey.config.WormConfiguration.PBA_linux_filename",
+        # Current server is used for attack telemetry
-        CUSTOM_LINUX_FILENAME,
+        "current_server": CUSTOM_SERVER,
-    )
+    }
-    return UsersPBA(MagicMock())
 
 
-def test_command_linux_custom_file_and_cmd(
+def test_command_linux_custom_file_and_cmd(fake_custom_pba_linux_options, set_os_linux):
-    mock_UsersPBA_linux_custom_file_and_cmd,
+    pba = CustomPBA(MagicMock())
-):
+    pba._set_options(fake_custom_pba_linux_options)
     expected_command = f"cd {MONKEY_DIR_PATH} ; {CUSTOM_LINUX_CMD}"
-    assert mock_UsersPBA_linux_custom_file_and_cmd.command == expected_command
+    assert pba.command == expected_command
+    assert pba.filename == CUSTOM_LINUX_FILENAME
 
 
 @pytest.fixture
-def mock_UsersPBA_windows_custom_file_and_cmd(set_os_windows, fake_monkey_dir_path, monkeypatch):
+def fake_custom_pba_windows_options():
-    monkeypatch.setattr(
+    return {
-        "infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd",
+        "linux_command": "",
-        CUSTOM_WINDOWS_CMD,
+        "linux_filename": "",
-    )
+        "windows_command": CUSTOM_WINDOWS_CMD,
-    monkeypatch.setattr(
+        "windows_filename": CUSTOM_WINDOWS_FILENAME,
-        "infection_monkey.config.WormConfiguration.PBA_windows_filename",
+        # Current server is used for attack telemetry
-        CUSTOM_WINDOWS_FILENAME,
+        "current_server": CUSTOM_SERVER,
-    )
+    }
-    return UsersPBA(MagicMock())
 
 
-def test_command_windows_custom_file_and_cmd(
+def test_command_windows_custom_file_and_cmd(fake_custom_pba_windows_options, set_os_windows):
-    mock_UsersPBA_windows_custom_file_and_cmd,
+
-):
+    pba = CustomPBA(MagicMock())
+    pba._set_options(fake_custom_pba_windows_options)
     expected_command = f"cd {MONKEY_DIR_PATH} & {CUSTOM_WINDOWS_CMD}"
-    assert mock_UsersPBA_windows_custom_file_and_cmd.command == expected_command
+    assert pba.command == expected_command
+    assert pba.filename == CUSTOM_WINDOWS_FILENAME
 
 
 @pytest.fixture
-def mock_UsersPBA_linux_custom_file(set_os_linux, fake_monkey_dir_path, monkeypatch):
+def fake_options_files_only():
-    monkeypatch.setattr("infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", None)
+    return {
-    monkeypatch.setattr(
+        "linux_command": "",
-        "infection_monkey.config.WormConfiguration.PBA_linux_filename",
+        "linux_filename": CUSTOM_LINUX_FILENAME,
-        CUSTOM_LINUX_FILENAME,
+        "windows_command": "",
-    )
+        "windows_filename": CUSTOM_WINDOWS_FILENAME,
-    return UsersPBA(MagicMock())
+        # Current server is used for attack telemetry
+        "current_server": CUSTOM_SERVER,
+    }
 
 
-def test_command_linux_custom_file(mock_UsersPBA_linux_custom_file):
+@pytest.mark.parametrize("os", [set_os_linux, set_os_windows])
-    expected_command = ""
+def test_files_only(fake_options_files_only, os):
-    assert mock_UsersPBA_linux_custom_file.command == expected_command
+    pba = CustomPBA(MagicMock())
+    pba._set_options(fake_options_files_only)
+    assert pba.command == ""
 
 
 @pytest.fixture
-def mock_UsersPBA_windows_custom_file(set_os_windows, fake_monkey_dir_path, monkeypatch):
+def fake_options_commands_only():
-    monkeypatch.setattr("infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", None)
+    return {
-    monkeypatch.setattr(
+        "linux_command": CUSTOM_LINUX_CMD,
-        "infection_monkey.config.WormConfiguration.PBA_windows_filename",
+        "linux_filename": "",
-        CUSTOM_WINDOWS_FILENAME,
+        "windows_command": CUSTOM_WINDOWS_CMD,
-    )
+        "windows_filename": "",
-    return UsersPBA(MagicMock())
+        # Current server is used for attack telemetry
+        "current_server": CUSTOM_SERVER,
+    }
 
 
-def test_command_windows_custom_file(mock_UsersPBA_windows_custom_file):
+def test_commands_only(fake_options_commands_only, set_os_linux):
-    expected_command = ""
+    pba = CustomPBA(MagicMock())
-    assert mock_UsersPBA_windows_custom_file.command == expected_command
+    pba._set_options(fake_options_commands_only)
+    assert pba.command == CUSTOM_LINUX_CMD
+    assert pba.filename == ""
 
 
-@pytest.fixture
+def test_commands_only_windows(fake_options_commands_only, set_os_windows):
-def mock_UsersPBA_linux_custom_cmd(set_os_linux, fake_monkey_dir_path, monkeypatch):
+    pba = CustomPBA(MagicMock())
-    monkeypatch.setattr(
+    pba._set_options(fake_options_commands_only)
-        "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd",
+    assert pba.command == CUSTOM_WINDOWS_CMD
-        CUSTOM_LINUX_CMD,
+    assert pba.filename == ""
-    )
-    monkeypatch.setattr("infection_monkey.config.WormConfiguration.PBA_linux_filename", None)
-    return UsersPBA(MagicMock())
-
-
-def test_command_linux_custom_cmd(mock_UsersPBA_linux_custom_cmd):
-    expected_command = CUSTOM_LINUX_CMD
-    assert mock_UsersPBA_linux_custom_cmd.command == expected_command
-
-
-@pytest.fixture
-def mock_UsersPBA_windows_custom_cmd(set_os_windows, fake_monkey_dir_path, monkeypatch):
-    monkeypatch.setattr(
-        "infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd",
-        CUSTOM_WINDOWS_CMD,
-    )
-    monkeypatch.setattr("infection_monkey.config.WormConfiguration.PBA_windows_filename", None)
-    return UsersPBA(MagicMock())
-
-
-def test_command_windows_custom_cmd(mock_UsersPBA_windows_custom_cmd):
-    expected_command = CUSTOM_WINDOWS_CMD
-    assert mock_UsersPBA_windows_custom_cmd.command == expected_command

monkey/tests/unit_tests/monkey_island/cc/services/test_config.py

@@ -66,12 +66,6 @@ def test_format_config_for_agent__pbas(flat_monkey_config):
         "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)
 
@@ -84,6 +78,19 @@ def test_format_config_for_agent__pbas(flat_monkey_config):
     assert "PBA_windows_filename" not in flat_monkey_config
 
 
+def test_format_config_for_custom_pbas(flat_monkey_config):
+    custom_config = {
+        "linux_command": "bash test.sh",
+        "windows_command": "powershell test.ps1",
+        "linux_filename": "test.sh",
+        "windows_filename": "test.ps1",
+        "current_server": "10.197.94.72:5000",
+    }
+    ConfigService.format_flat_config_for_agent(flat_monkey_config)
+
+    assert flat_monkey_config["custom_pbas"] == custom_config
+
+
 def test_get_config_propagation_credentials_from_flat_config(flat_monkey_config):
     expected_creds = {
         "exploit_lm_hash_list": ["lm_hash_1", "lm_hash_2"],