Merge pull request #1829 from guardicore/1604-modify-puppet-to-run-pbas

PBA modifications so they run

monkey/infection_monkey/master/automated_master.py

@@ -201,9 +201,7 @@ class AutomatedMaster(IMaster):
             return
 
         for pba_data in self._puppet.run_pba(name, options):
-            self._telemetry_messenger.send_telemetry(
+            self._telemetry_messenger.send_telemetry(PostBreachTelem(pba_data))
-                PostBreachTelem(pba_data.display_name, pba_data.command, pba_data.result)
-            )
 
     def _can_propagate(self) -> bool:
         return True

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,12 +14,13 @@ 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
             self.pba_data.append(PostBreachData(self.name, self.command, results))
-            return self.pba_data
+
+        return self.pba_data
 
     def clear_command_history_pba_list(self):
         return self.CommandHistoryPBAGenerator().get_clear_command_history_pbas()
@@ -44,7 +46,11 @@ class ClearCommandHistory(PBA):
 
         class ClearCommandHistoryFile(PBA):
             def __init__(self, linux_cmds):
-                super().__init__(name=POST_BREACH_CLEAR_CMD_HISTORY, linux_cmd=linux_cmds)
+                super().__init__(
+                    telemetry_messenger=None,
+                    name=POST_BREACH_CLEAR_CMD_HISTORY,
+                    linux_cmd=linux_cmds,
+                )
 
             def run(self):
                 if self.command:

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,16 +23,17 @@ 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()
             super(HiddenFiles, self).__init__(
+                self.telemetry_messenger,
                 name=POST_BREACH_HIDDEN_FILES,
                 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 = [
@@ -63,6 +64,7 @@ class ModifyShellStartupFiles(PBA):
         class ModifyShellStartupFile(PBA):
             def __init__(self, linux_cmds, windows_cmds):
                 super().__init__(
+                    telemetry_messenger=None,
                     name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION,
                     linux_cmd=linux_cmds,
                     windows_cmd=windows_cmds,

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/post_breach/pba.py

@@ -1,6 +1,6 @@
 import logging
 import subprocess
-from typing import Iterable
+from typing import Dict, Iterable
 
 from common.utils.attack_utils import ScanStatus
 from infection_monkey.i_puppet.i_puppet import PostBreachData
@@ -30,7 +30,7 @@ class PBA:
         self.pba_data = []
         self.telemetry_messenger = telemetry_messenger
 
-    def run(self) -> Iterable[PostBreachData]:
+    def run(self, options: Dict) -> Iterable[PostBreachData]:
         """
         Runs post breach action command
         """
@@ -45,10 +45,11 @@ class PBA:
                     )
                 )
             self.pba_data.append(PostBreachData(self.name, self.command, result))
-            return self.pba_data
         else:
             logger.debug(f"No command available for PBA '{self.name}' on current OS, skipping.")
 
+        return self.pba_data
+
     def is_script(self):
         """
         Determines if PBA is a script (PBA might be a single command)

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)