Agent: Change PBA constructor to accept telemetry messenger

This change allows to run different PBA's with different telemetry messengers

monkey/infection_monkey/monkey.py

@@ -248,26 +248,56 @@ class InfectionMonkey:
         )
 
         puppet.load_plugin(
-            "CommunicateAsBackdoorUser", CommunicateAsBackdoorUser, PluginType.POST_BREACH_ACTION
+            "CommunicateAsBackdoorUser",
+            CommunicateAsBackdoorUser(self.telemetry_messenger),
+            PluginType.POST_BREACH_ACTION,
         )
         puppet.load_plugin(
-            "ModifyShellStartupFiles", ModifyShellStartupFiles, PluginType.POST_BREACH_ACTION
-        )
-        puppet.load_plugin("HiddenFiles", HiddenFiles, PluginType.POST_BREACH_ACTION)
-        puppet.load_plugin("TrapCommand", CommunicateAsBackdoorUser, PluginType.POST_BREACH_ACTION)
-        puppet.load_plugin("ChangeSetuidSetgid", ChangeSetuidSetgid, PluginType.POST_BREACH_ACTION)
-        puppet.load_plugin("ScheduleJobs", ScheduleJobs, PluginType.POST_BREACH_ACTION)
-        puppet.load_plugin("Timestomping", Timestomping, PluginType.POST_BREACH_ACTION)
-        puppet.load_plugin("AccountDiscovery", AccountDiscovery, PluginType.POST_BREACH_ACTION)
-        puppet.load_plugin(
-            "ProcessListCollection", ProcessListCollection, PluginType.POST_BREACH_ACTION
-        )
-        puppet.load_plugin("TrapCommand", TrapCommand, PluginType.POST_BREACH_ACTION)
-        puppet.load_plugin(
-            "SignedScriptProxyExecution", SignedScriptProxyExecution, PluginType.POST_BREACH_ACTION
+            "ModifyShellStartupFiles",
+            ModifyShellStartupFiles(self.telemetry_messenger),
+            PluginType.POST_BREACH_ACTION,
         )
         puppet.load_plugin(
-            "ClearCommandHistory", ClearCommandHistory, PluginType.POST_BREACH_ACTION
+            "HiddenFiles", HiddenFiles(self.telemetry_messenger), PluginType.POST_BREACH_ACTION
+        )
+        puppet.load_plugin(
+            "TrapCommand",
+            CommunicateAsBackdoorUser(self.telemetry_messenger),
+            PluginType.POST_BREACH_ACTION,
+        )
+        puppet.load_plugin(
+            "ChangeSetuidSetgid",
+            ChangeSetuidSetgid(self.telemetry_messenger),
+            PluginType.POST_BREACH_ACTION,
+        )
+        puppet.load_plugin(
+            "ScheduleJobs", ScheduleJobs(self.telemetry_messenger), PluginType.POST_BREACH_ACTION
+        )
+        puppet.load_plugin(
+            "Timestomping", Timestomping(self.telemetry_messenger), PluginType.POST_BREACH_ACTION
+        )
+        puppet.load_plugin(
+            "AccountDiscovery",
+            AccountDiscovery(self.telemetry_messenger),
+            PluginType.POST_BREACH_ACTION,
+        )
+        puppet.load_plugin(
+            "ProcessListCollection",
+            ProcessListCollection(self.telemetry_messenger),
+            PluginType.POST_BREACH_ACTION,
+        )
+        puppet.load_plugin(
+            "TrapCommand", TrapCommand(self.telemetry_messenger), PluginType.POST_BREACH_ACTION
+        )
+        puppet.load_plugin(
+            "SignedScriptProxyExecution",
+            SignedScriptProxyExecution(self.telemetry_messenger),
+            PluginType.POST_BREACH_ACTION,
+        )
+        puppet.load_plugin(
+            "ClearCommandHistory",
+            ClearCommandHistory(self.telemetry_messenger),
+            PluginType.POST_BREACH_ACTION,
         )
 
         puppet.load_plugin("ransomware", RansomwarePayload(), PluginType.PAYLOAD)

monkey/infection_monkey/post_breach/actions/change_file_privileges.py

@@ -3,11 +3,12 @@ from infection_monkey.post_breach.pba import PBA
 from infection_monkey.post_breach.setuid_setgid.setuid_setgid import (
     get_commands_to_change_setuid_setgid,
 )
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
 
 class ChangeSetuidSetgid(PBA):
-    def __init__(self):
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
         linux_cmds = get_commands_to_change_setuid_setgid()
         super(ChangeSetuidSetgid, self).__init__(
-            POST_BREACH_SETUID_SETGID, linux_cmd=" ".join(linux_cmds)
+            telemetry_messenger, POST_BREACH_SETUID_SETGID, linux_cmd=" ".join(linux_cmds)
         )

monkey/infection_monkey/post_breach/actions/clear_command_history.py

@@ -6,20 +6,21 @@ from infection_monkey.post_breach.clear_command_history.clear_command_history im
     get_commands_to_clear_command_history,
 )
 from infection_monkey.post_breach.pba import PBA
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
 
 class ClearCommandHistory(PBA):
-    def __init__(self):
-        super().__init__(name=POST_BREACH_CLEAR_CMD_HISTORY)
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
+        super().__init__(telemetry_messenger, name=POST_BREACH_CLEAR_CMD_HISTORY)
 
     def run(self):
-        results = [pba.run() for pba in self.clear_command_history_PBA_list()]
+        results = [pba.run() for pba in self.clear_command_history_pba_list()]
         if results:
             # `self.command` is empty here
             self.pba_data.append(PostBreachData(self.name, self.command, results))
             return self.pba_data
 
-    def clear_command_history_PBA_list(self):
+    def clear_command_history_pba_list(self):
         return self.CommandHistoryPBAGenerator().get_clear_command_history_pbas()
 
     class CommandHistoryPBAGenerator:

monkey/infection_monkey/post_breach/actions/collect_processes_list.py

@@ -5,6 +5,7 @@ import psutil
 from common.common_consts.post_breach_consts import POST_BREACH_PROCESS_LIST_COLLECTION
 from infection_monkey.i_puppet.i_puppet import PostBreachData
 from infection_monkey.post_breach.pba import PBA
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
 logger = logging.getLogger(__name__)
 
@@ -17,8 +18,8 @@ except NameError:
 
 
 class ProcessListCollection(PBA):
-    def __init__(self):
-        super().__init__(POST_BREACH_PROCESS_LIST_COLLECTION)
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
+        super().__init__(telemetry_messenger, POST_BREACH_PROCESS_LIST_COLLECTION)
 
     def run(self):
         """

monkey/infection_monkey/post_breach/actions/communicate_as_backdoor_user.py

@@ -7,6 +7,7 @@ import subprocess
 from common.common_consts.post_breach_consts import POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER
 from infection_monkey.i_puppet.i_puppet import PostBreachData
 from infection_monkey.post_breach.pba import PBA
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.utils.auto_new_user_factory import create_auto_new_user
 from infection_monkey.utils.environment import is_windows_os
 from infection_monkey.utils.new_user_error import NewUserError
@@ -33,9 +34,9 @@ class CommunicateAsBackdoorUser(PBA):
     are created.
     """
 
-    def __init__(self):
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
         super(CommunicateAsBackdoorUser, self).__init__(
-            name=POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER
+            telemetry_messenger, name=POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER
         )
 
     def run(self):

monkey/infection_monkey/post_breach/actions/discover_accounts.py

@@ -3,11 +3,15 @@ from infection_monkey.post_breach.account_discovery.account_discovery import (
     get_commands_to_discover_accounts,
 )
 from infection_monkey.post_breach.pba import PBA
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
 
 class AccountDiscovery(PBA):
-    def __init__(self):
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
         linux_cmds, windows_cmds = get_commands_to_discover_accounts()
         super().__init__(
-            POST_BREACH_ACCOUNT_DISCOVERY, linux_cmd=" ".join(linux_cmds), windows_cmd=windows_cmds
+            telemetry_messenger,
+            POST_BREACH_ACCOUNT_DISCOVERY,
+            linux_cmd=" ".join(linux_cmds),
+            windows_cmd=windows_cmds,
         )

monkey/infection_monkey/post_breach/actions/hide_files.py

@@ -1,6 +1,7 @@
 from common.common_consts.post_breach_consts import POST_BREACH_HIDDEN_FILES
 from infection_monkey.i_puppet.i_puppet import PostBreachData
 from infection_monkey.post_breach.pba import PBA
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.utils.environment import is_windows_os
 from infection_monkey.utils.hidden_files import (
     cleanup_hidden_files,
@@ -17,8 +18,8 @@ class HiddenFiles(PBA):
     This PBA attempts to create hidden files and folders.
     """
 
-    def __init__(self):
-        super(HiddenFiles, self).__init__(name=POST_BREACH_HIDDEN_FILES)
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
+        super(HiddenFiles, self).__init__(telemetry_messenger, name=POST_BREACH_HIDDEN_FILES)
 
     def run(self):
         # create hidden files and folders

monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py

@@ -6,6 +6,7 @@ from infection_monkey.post_breach.pba import PBA
 from infection_monkey.post_breach.shell_startup_files.shell_startup_files_modification import (
     get_commands_to_modify_shell_startup_files,
 )
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
 
 class ModifyShellStartupFiles(PBA):
@@ -15,8 +16,8 @@ class ModifyShellStartupFiles(PBA):
     and profile.ps1 in windows.
     """
 
-    def __init__(self):
-        super().__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION)
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
+        super().__init__(telemetry_messenger, name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION)
 
     def run(self):
         results = [pba.run() for pba in self.modify_shell_startup_PBA_list()]

monkey/infection_monkey/post_breach/actions/schedule_jobs.py

@@ -4,6 +4,7 @@ from infection_monkey.post_breach.job_scheduling.job_scheduling import (
     remove_scheduled_jobs,
 )
 from infection_monkey.post_breach.pba import PBA
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
 
 class ScheduleJobs(PBA):
@@ -11,10 +12,11 @@ class ScheduleJobs(PBA):
     This PBA attempts to schedule jobs on the system.
     """
 
-    def __init__(self):
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
         linux_cmds, windows_cmds = get_commands_to_schedule_jobs()
 
         super(ScheduleJobs, self).__init__(
+            telemetry_messenger,
             name=POST_BREACH_JOB_SCHEDULING,
             linux_cmd=" ".join(linux_cmds),
             windows_cmd=windows_cmds,

monkey/infection_monkey/post_breach/actions/timestomping.py

@@ -1,9 +1,15 @@
 from common.common_consts.post_breach_consts import POST_BREACH_TIMESTOMPING
 from infection_monkey.post_breach.pba import PBA
 from infection_monkey.post_breach.timestomping.timestomping import get_timestomping_commands
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
 
 class Timestomping(PBA):
-    def __init__(self):
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
         linux_cmds, windows_cmds = get_timestomping_commands()
-        super().__init__(POST_BREACH_TIMESTOMPING, linux_cmd=linux_cmds, windows_cmd=windows_cmds)
+        super().__init__(
+            telemetry_messenger,
+            POST_BREACH_TIMESTOMPING,
+            linux_cmd=linux_cmds,
+            windows_cmd=windows_cmds,
+        )

monkey/infection_monkey/post_breach/actions/use_signed_scripts.py

@@ -7,15 +7,20 @@ from infection_monkey.post_breach.signed_script_proxy.signed_script_proxy import
     cleanup_changes,
     get_commands_to_proxy_execution_using_signed_script,
 )
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.utils.environment import is_windows_os
 
 logger = logging.getLogger(__name__)
 
 
 class SignedScriptProxyExecution(PBA):
-    def __init__(self):
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
         windows_cmds = get_commands_to_proxy_execution_using_signed_script()
-        super().__init__(POST_BREACH_SIGNED_SCRIPT_PROXY_EXEC, windows_cmd=" ".join(windows_cmds))
+        super().__init__(
+            telemetry_messenger,
+            POST_BREACH_SIGNED_SCRIPT_PROXY_EXEC,
+            windows_cmd=" ".join(windows_cmds),
+        )
 
     def run(self):
         original_comspec = ""

monkey/infection_monkey/post_breach/actions/use_trap_command.py

@@ -1,9 +1,12 @@
 from common.common_consts.post_breach_consts import POST_BREACH_TRAP_COMMAND
 from infection_monkey.post_breach.pba import PBA
 from infection_monkey.post_breach.trap_command.trap_command import get_trap_commands
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
 
 class TrapCommand(PBA):
-    def __init__(self):
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
         linux_cmds = get_trap_commands()
-        super(TrapCommand, self).__init__(POST_BREACH_TRAP_COMMAND, linux_cmd=" ".join(linux_cmds))
+        super(TrapCommand, self).__init__(
+            telemetry_messenger, POST_BREACH_TRAP_COMMAND, linux_cmd=" ".join(linux_cmds)
+        )

monkey/infection_monkey/post_breach/actions/users_custom_pba.py

@@ -8,6 +8,7 @@ from infection_monkey.control import ControlClient
 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
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.utils.environment import is_windows_os
 from infection_monkey.utils.monkey_dir import get_monkey_dir_path
 
@@ -23,8 +24,8 @@ class UsersPBA(PBA):
     Defines user's configured post breach action.
     """
 
-    def __init__(self):
-        super(UsersPBA, self).__init__(POST_BREACH_FILE_EXECUTION)
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
+        super(UsersPBA, self).__init__(telemetry_messenger, POST_BREACH_FILE_EXECUTION)
         self.filename = ""
 
         if not is_windows_os():

monkey/infection_monkey/post_breach/pba.py

@@ -6,7 +6,6 @@ from common.utils.attack_utils import ScanStatus
 from infection_monkey.i_puppet.i_puppet import PostBreachData
 from infection_monkey.telemetry.attack.t1064_telem import T1064Telem
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
-from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
 from infection_monkey.utils.environment import is_windows_os
 
 logger = logging.getLogger(__name__)
@@ -18,7 +17,9 @@ class PBA:
     machine.
     """
 
-    def __init__(self, name="unknown", linux_cmd="", windows_cmd=""):
+    def __init__(
+        self, telemetry_messenger: ITelemetryMessenger, name="unknown", linux_cmd="", windows_cmd=""
+    ):
         """
         :param name: Name of post breach action.
         :param linux_cmd: Command that will be executed on breached machine
@@ -27,6 +28,7 @@ class PBA:
         self.command = PBA.choose_command(linux_cmd, windows_cmd)
         self.name = name
         self.pba_data = []
+        self.telemetry_messenger = telemetry_messenger
 
     def run(self) -> Iterable[PostBreachData]:
         """