Agent: Modify puppet to run PBAs instead of using the mock puppet

monkey/infection_monkey/post_breach/actions/clear_command_history.py

@@ -1,4 +1,5 @@
 import subprocess
+from typing import Dict
 
 from common.common_consts.post_breach_consts import POST_BREACH_CLEAR_CMD_HISTORY
 from infection_monkey.i_puppet.i_puppet import PostBreachData
@@ -13,7 +14,7 @@ class ClearCommandHistory(PBA):
     def __init__(self, telemetry_messenger: ITelemetryMessenger):
         super().__init__(telemetry_messenger, name=POST_BREACH_CLEAR_CMD_HISTORY)
 
-    def run(self):
+    def run(self, options: Dict):
         results = [pba.run() for pba in self.clear_command_history_pba_list()]
         if results:
             # `self.command` is empty here

monkey/infection_monkey/post_breach/actions/collect_processes_list.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Dict
 
 import psutil
 
@@ -21,7 +22,7 @@ class ProcessListCollection(PBA):
     def __init__(self, telemetry_messenger: ITelemetryMessenger):
         super().__init__(telemetry_messenger, POST_BREACH_PROCESS_LIST_COLLECTION)
 
-    def run(self):
+    def run(self, options: Dict):
         """
         Collects process information from the host.
         Currently lists process name, ID, parent ID, command line

monkey/infection_monkey/post_breach/actions/communicate_as_backdoor_user.py

@@ -3,6 +3,7 @@ import random
 import shutil
 import string
 import subprocess
+from typing import Dict
 
 from common.common_consts.post_breach_consts import POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER
 from infection_monkey.i_puppet.i_puppet import PostBreachData
@@ -39,7 +40,7 @@ class CommunicateAsBackdoorUser(PBA):
             telemetry_messenger, name=POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER
         )
 
-    def run(self):
+    def run(self, options: Dict):
         username = CommunicateAsBackdoorUser.get_random_new_user_name()
         try:
             password = get_random_password(14)

monkey/infection_monkey/post_breach/actions/hide_files.py

@@ -1,3 +1,5 @@
+from typing import Dict
+
 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
@@ -21,7 +23,7 @@ class HiddenFiles(PBA):
     def __init__(self, telemetry_messenger: ITelemetryMessenger):
         super(HiddenFiles, self).__init__(telemetry_messenger, name=POST_BREACH_HIDDEN_FILES)
 
-    def run(self):
+    def run(self, options: Dict):
         # create hidden files and folders
         for function_to_get_commands in HIDDEN_FSO_CREATION_COMMANDS:
             linux_cmds, windows_cmds = function_to_get_commands()
@@ -30,7 +32,7 @@ class HiddenFiles(PBA):
                 linux_cmd=" ".join(linux_cmds),
                 windows_cmd=windows_cmds,
             )
-            super(HiddenFiles, self).run()
+            super(HiddenFiles, self).run(options)
 
         if is_windows_os():  # use winAPI
             result, status = get_winAPI_to_hide_files()

monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py

@@ -1,4 +1,5 @@
 import subprocess
+from typing import Dict
 
 from common.common_consts.post_breach_consts import POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION
 from infection_monkey.i_puppet.i_puppet import PostBreachData
@@ -19,7 +20,7 @@ class ModifyShellStartupFiles(PBA):
     def __init__(self, telemetry_messenger: ITelemetryMessenger):
         super().__init__(telemetry_messenger, name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION)
 
-    def run(self):
+    def run(self, options: Dict):
         results = [pba.run() for pba in self.modify_shell_startup_PBA_list()]
         if not results:
             results = [

monkey/infection_monkey/post_breach/actions/schedule_jobs.py

@@ -1,3 +1,5 @@
+from typing import Dict
+
 from common.common_consts.post_breach_consts import POST_BREACH_JOB_SCHEDULING
 from infection_monkey.post_breach.job_scheduling.job_scheduling import (
     get_commands_to_schedule_jobs,
@@ -22,7 +24,7 @@ class ScheduleJobs(PBA):
             windows_cmd=windows_cmds,
         )
 
-    def run(self):
+    def run(self, options: Dict):
-        super(ScheduleJobs, self).run()
+        super(ScheduleJobs, self).run(options)
         remove_scheduled_jobs()
         return self.pba_data

monkey/infection_monkey/post_breach/actions/use_signed_scripts.py

@@ -1,5 +1,6 @@
 import logging
 import subprocess
+from typing import Dict
 
 from common.common_consts.post_breach_consts import POST_BREACH_SIGNED_SCRIPT_PROXY_EXEC
 from infection_monkey.post_breach.pba import PBA
@@ -22,14 +23,14 @@ class SignedScriptProxyExecution(PBA):
             windows_cmd=" ".join(windows_cmds),
         )
 
-    def run(self):
+    def run(self, options: Dict):
         original_comspec = ""
         try:
             if is_windows_os():
                 original_comspec = subprocess.check_output(  # noqa: DUO116
                     "if defined COMSPEC echo %COMSPEC%", shell=True
                 ).decode()
-            super().run()
+            super().run(options)
             return self.pba_data
         except Exception as e:
             logger.warning(

monkey/infection_monkey/puppet/puppet.py

@@ -37,7 +37,8 @@ class Puppet(IPuppet):
         return credential_collector.collect_credentials(options)
 
     def run_pba(self, name: str, options: Dict) -> Iterable[PostBreachData]:
-        return self._mock_puppet.run_pba(name, options)
+        pba = self._plugin_registry.get_plugin(name, PluginType.POST_BREACH_ACTION)
+        return pba.run(options)
 
     def ping(self, host: str, timeout: float = CONNECTION_TIMEOUT) -> PingScanData:
         return network_scanning.ping(host, timeout)