Agent: Load and fix the custom PBA into puppet

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 Any, Dict, Iterable, List, Mapping, Sequence
+from typing import Dict, Iterable, List, Mapping, Sequence
 
 from infection_monkey.model import VictimHost
 
@@ -67,14 +67,6 @@ 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

@@ -4,7 +4,6 @@ import time
 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
@@ -155,7 +154,7 @@ class AutomatedMaster(IMaster):
             ),
         )
         pba_thread = create_daemon_thread(
-            target=self._run_PBAs,
+            target=self._run_pbas,
             name="PBAThread",
             args=(config["post_breach_actions"].items(), self._run_pba, config["custom_pbas"]),
         )
@@ -197,10 +196,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))
 
@@ -213,14 +208,12 @@ class AutomatedMaster(IMaster):
 
         self._puppet.run_payload(name, options, self._stop)
 
-    def _run_PBAs(
+    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)
+        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/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.users_custom_pba import UsersPBA
 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", UsersPBA(self._telemetry_messenger), PluginType.POST_BREACH_ACTION
+        )
 
         puppet.load_plugin("ransomware", RansomwarePayload(), PluginType.PAYLOAD)
 

monkey/infection_monkey/post_breach/custom_pba/users_custom_pba.py

@@ -1,11 +1,11 @@
 import logging
 import os
-from typing import Any, Mapping
+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
@@ -25,10 +25,18 @@ class UsersPBA(PBA):
     Defines user's configured post breach action.
     """
 
-    def __init__(self, options: Mapping[str, Any], telemetry_messenger: ITelemetryMessenger):
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
         super(UsersPBA, self).__init__(telemetry_messenger, POST_BREACH_FILE_EXECUTION)
         self.filename = ""
 
+    def run(self, options: Dict) -> Iterable[PostBreachData]:
+        self._set_options(options)
+        return super().run(options)
+
+    def _set_options(self, options: Dict):
+        # Required for attack telemetry
+        self.current_server = options["current_server"]
+
         if is_windows_os():
             # Add windows commands to PBA's
             if options["windows_filename"]:
@@ -54,7 +62,7 @@ class UsersPBA(PBA):
 
     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()
 
     @staticmethod
@@ -85,11 +93,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],
-                get_interface_to_target(WormConfiguration.current_server.split(":")[0]),
+                self.current_server.split(":")[0],
+                get_interface_to_target(self.current_server.split(":")[0]),
                 filename,
             )
         )

monkey/monkey_island/cc/services/config.py

@@ -457,6 +457,8 @@ class ConfigService:
             "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/master/mock_puppet.py

@@ -1,6 +1,6 @@
 import logging
 import threading
-from typing import Any, Dict, Iterable, List, Mapping, Sequence
+from typing import Dict, Iterable, List, Sequence
 
 from infection_monkey.credential_collectors import LMHash, Password, SSHKeypair, Username
 from infection_monkey.i_puppet import (
@@ -57,9 +57,6 @@ 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: