Island, Agent: Add custom user PBA to puppet and master

monkey/infection_monkey/i_puppet/i_puppet.py

@@ -3,7 +3,7 @@ import threading
 from collections import namedtuple
 from dataclasses import dataclass
 from enum import Enum
-from typing import Dict, Iterable, List, Mapping, Sequence
+from typing import Any, Dict, Iterable, List, Mapping, Sequence
 
 from infection_monkey.model import VictimHost
 
@@ -67,6 +67,14 @@ class IPuppet(metaclass=abc.ABCMeta):
         :rtype: Iterable[PostBreachData]
         """
 
+    @abc.abstractmethod
+    def run_custom_pba(self, options: Mapping[str, Any]) -> PostBreachData:
+        """
+        Runs a user configured post breach action (PBA)
+        :param Dict options: A dictionary containing options that modify the behavior of the PBA
+        :rtype: PostBreachData
+        """
+
     @abc.abstractmethod
     def ping(self, host: str, timeout: float) -> PingScanData:
         """

monkey/infection_monkey/master/automated_master.py

@@ -1,9 +1,10 @@
 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 common.common_consts.post_breach_consts import POST_BREACH_FILE_EXECUTION
 from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
 from infection_monkey.i_master import IMaster
 from infection_monkey.i_puppet import IPuppet
@@ -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()
@@ -212,6 +213,15 @@ 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)
+
+        command, result = self._puppet.run_custom_pba(custom_pba_options)
+        telem = PostBreachTelem(POST_BREACH_FILE_EXECUTION, command, result)
+        self._telemetry_messenger.send_telemetry(telem)
+
     def _run_plugins(
         self, plugins: Iterable[Any], plugin_type: str, callback: Callable[[Any], None]
     ):

monkey/infection_monkey/post_breach/custom_pba/users_custom_pba.py

@@ -1,5 +1,6 @@
 import logging
 import os
+from typing import Any, Mapping
 
 from common.common_consts.post_breach_consts import POST_BREACH_FILE_EXECUTION
 from common.utils.attack_utils import ScanStatus
@@ -24,32 +25,32 @@ class UsersPBA(PBA):
     Defines user's configured post breach action.
     """
 
-    def __init__(self, telemetry_messenger: ITelemetryMessenger):
+    def __init__(self, options: Mapping[str, Any], telemetry_messenger: ITelemetryMessenger):
         super(UsersPBA, self).__init__(telemetry_messenger, POST_BREACH_FILE_EXECUTION)
         self.filename = ""
 
-        if not is_windows_os():
-            # Add linux commands to PBA's
-            if WormConfiguration.PBA_linux_filename:
-                self.filename = WormConfiguration.PBA_linux_filename
-                if WormConfiguration.custom_PBA_linux_cmd:
-                    # Add change dir command, because user will try to access his file
-                    self.command = (
-                        DIR_CHANGE_LINUX % get_monkey_dir_path()
-                    ) + WormConfiguration.custom_PBA_linux_cmd
-            elif WormConfiguration.custom_PBA_linux_cmd:
-                self.command = WormConfiguration.custom_PBA_linux_cmd
-        else:
+        if is_windows_os():
             # Add windows commands to PBA's
-            if WormConfiguration.PBA_windows_filename:
-                self.filename = WormConfiguration.PBA_windows_filename
-                if WormConfiguration.custom_PBA_windows_cmd:
+            if options["windows_filename"]:
+                self.filename = options["windows_filename"]
+                if options["windows_command"]:
                     # Add change dir command, because user will try to access his file
-                    self.command = (
-                        DIR_CHANGE_WINDOWS % get_monkey_dir_path()
-                    ) + WormConfiguration.custom_PBA_windows_cmd
-            elif WormConfiguration.custom_PBA_windows_cmd:
-                self.command = WormConfiguration.custom_PBA_windows_cmd
+                    self.command = (DIR_CHANGE_WINDOWS % get_monkey_dir_path()) + options[
+                        "windows_command"
+                    ]
+            elif options["windows_command"]:
+                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:
@@ -57,12 +58,12 @@ class UsersPBA(PBA):
         return super(UsersPBA, self)._execute_default()
 
     @staticmethod
-    def should_run(class_name):
+    def should_run(options):
         if not is_windows_os():
-            if WormConfiguration.PBA_linux_filename or WormConfiguration.custom_PBA_linux_cmd:
+            if options["linux_filename"] or options["linux_command"]:
                 return True
         else:
-            if WormConfiguration.PBA_windows_filename or WormConfiguration.custom_PBA_windows_cmd:
+            if options["windows_filename"] or options["windows_command"]:
                 return True
         return False
 

monkey/monkey_island/cc/services/config.py

@@ -452,7 +452,7 @@ 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, ""),

monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py

@@ -1,6 +1,6 @@
 import logging
 import threading
-from typing import Dict, Iterable, List, Sequence
+from typing import Any, Dict, Iterable, List, Mapping, Sequence
 
 from infection_monkey.credential_collectors import LMHash, Password, SSHKeypair, Username
 from infection_monkey.i_puppet import (
@@ -57,6 +57,9 @@ class MockPuppet(IPuppet):
         else:
             return [PostBreachData(name, "pba command 2", ["pba result 2", False])]
 
+    def run_custom_pba(self, options: Mapping[str, Any]) -> PostBreachData:
+        pass
+
     def ping(self, host: str, timeout: float = 1) -> PingScanData:
         logger.debug(f"run_ping({host}, {timeout})")
         if host == DOT_1:

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

@@ -2,7 +2,7 @@ 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.users_custom_pba import UsersPBA
 
 MONKEY_DIR_PATH = "/dir/to/monkey/"
 CUSTOM_LINUX_CMD = "command-for-linux"