Merge pull request #1822 from guardicore/1604-modify-pbas-to-return-postbreachdata

Modify PBAs to return PostBreachData

monkey/infection_monkey/master/automated_master.py

@@ -198,8 +198,14 @@ class AutomatedMaster(IMaster):
         name = pba[0]
         options = pba[1]
 
-        result = self._puppet.run_pba(name, options)
+        # TEMPORARY; TO AVOID ERRORS SINCE THIS ISN'T IMPLEMENTED YET
-        self._telemetry_messenger.send_telemetry(PostBreachTelem(result))
+        if name == "Custom":
+            return
+
+        for pba_data in self._puppet.run_pba(name, options):
+            self._telemetry_messenger.send_telemetry(
+                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,11 +1,11 @@
 import subprocess
 
 from common.common_consts.post_breach_consts import POST_BREACH_CLEAR_CMD_HISTORY
+from infection_monkey.i_puppet.i_puppet import PostBreachData
 from infection_monkey.post_breach.clear_command_history.clear_command_history import (
     get_commands_to_clear_command_history,
 )
 from infection_monkey.post_breach.pba import PBA
-from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
 
 
 class ClearCommandHistory(PBA):
@@ -15,7 +15,9 @@ class ClearCommandHistory(PBA):
     def run(self):
         results = [pba.run() for pba in self.clear_command_history_PBA_list()]
         if results:
-            PostBreachTelem(self, results).send()
+            # `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):
         return self.CommandHistoryPBAGenerator().get_clear_command_history_pbas()

monkey/infection_monkey/post_breach/actions/collect_processes_list.py

@@ -3,6 +3,7 @@ import logging
 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
 
 logger = logging.getLogger(__name__)
@@ -16,9 +17,6 @@ except NameError:
 
 
 class ProcessListCollection(PBA):
-    # TODO: (?) Move all PBA consts into their classes
-    display_name = POST_BREACH_PROCESS_LIST_COLLECTION
-
     def __init__(self):
         super().__init__(POST_BREACH_PROCESS_LIST_COLLECTION)
 
@@ -54,4 +52,6 @@ class ProcessListCollection(PBA):
                 }
                 continue
 
-        return self.command, (processes, success_state)
+        # No command here; used psutil
+        self.pba_data.append(PostBreachData(self.name, self.command, (processes, success_state)))
+        return self.pba_data

monkey/infection_monkey/post_breach/actions/communicate_as_backdoor_user.py

@@ -5,8 +5,8 @@ import string
 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.post_breach_telem import PostBreachTelem
 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
@@ -49,11 +49,20 @@ class CommunicateAsBackdoorUser(PBA):
                     )
                 )
                 exit_status = new_user.run_as(http_request_commandline)
-                self.send_result_telemetry(exit_status, http_request_commandline, username)
+                result = CommunicateAsBackdoorUser._get_result_for_telemetry(
+                    exit_status, http_request_commandline, username
+                )
+                # `command` is empty here; we could get the command from `new_user` but that
+                # doesn't work either since Windows doesn't use a command, it uses win32 modules
+                self.pba_data.append(PostBreachData(self.name, self.command, result))
         except subprocess.CalledProcessError as e:
-            PostBreachTelem(self, (e.output.decode(), False)).send()
+            self.pba_data.append(
+                PostBreachData(self.name, self.command, (e.output.decode(), False))
+            )
         except NewUserError as e:
-            PostBreachTelem(self, (str(e), False)).send()
+            self.pba_data.append(PostBreachData(self.name, self.command, (str(e), False)))
+        finally:
+            return self.pba_data
 
     @staticmethod
     def get_random_new_user_name():
@@ -79,28 +88,19 @@ class CommunicateAsBackdoorUser(PBA):
                 format_string = "wget -O/dev/null -q {url} --method=HEAD --timeout=10"
         return format_string.format(url=url)
 
-    def send_result_telemetry(self, exit_status, commandline, username):
+    @staticmethod
-        """
+    def _get_result_for_telemetry(exit_status, commandline, username):
-        Parses the result of the command and sends telemetry accordingly.
-
-        :param exit_status: In both Windows and Linux, 0 exit code indicates success.
-        :param commandline: Exact commandline which was executed, for reporting back.
-        :param username: Username from which the command was executed, for reporting back.
-        """
         if exit_status == 0:
-            PostBreachTelem(
+            result = (CREATED_PROCESS_AS_USER_SUCCESS_FORMAT.format(commandline, username), True)
-                self, (CREATED_PROCESS_AS_USER_SUCCESS_FORMAT.format(commandline, username), True)
-            ).send()
         else:
-            PostBreachTelem(
+            result = (
-                self,
+                CREATED_PROCESS_AS_USER_FAILED_FORMAT.format(
-                (
+                    commandline, username, exit_status, twos_complement(exit_status)
-                    CREATED_PROCESS_AS_USER_FAILED_FORMAT.format(
-                        commandline, username, exit_status, twos_complement(exit_status)
-                    ),
-                    False,
                 ),
-            ).send()
+                False,
+            )
+
+        return result
 
 
 def twos_complement(exit_status):

monkey/infection_monkey/post_breach/actions/hide_files.py

@@ -1,6 +1,6 @@
 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.post_breach_telem import PostBreachTelem
 from infection_monkey.utils.environment import is_windows_os
 from infection_monkey.utils.hidden_files import (
     cleanup_hidden_files,
@@ -30,9 +30,13 @@ class HiddenFiles(PBA):
                 windows_cmd=windows_cmds,
             )
             super(HiddenFiles, self).run()
+
         if is_windows_os():  # use winAPI
             result, status = get_winAPI_to_hide_files()
-            PostBreachTelem(self, (result, status)).send()
+            # no command here, used WinAPI
+            self.pba_data.append(PostBreachData(self.name, self.command, (result, status)))
 
         # cleanup hidden files and folders
         cleanup_hidden_files(is_windows_os())
+
+        return self.pba_data

monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py

@@ -1,11 +1,11 @@
 import subprocess
 
 from common.common_consts.post_breach_consts import POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION
+from infection_monkey.i_puppet.i_puppet import PostBreachData
 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.post_breach_telem import PostBreachTelem
 
 
 class ModifyShellStartupFiles(PBA):
@@ -27,13 +27,18 @@ class ModifyShellStartupFiles(PBA):
                     False,
                 )
             ]
-        PostBreachTelem(self, results).send()
+        # `command` is empty here since multiple commands were run through objects of the nested
+        # class. The results of each of those were aggregated to send the telemetry just once.
+        self.pba_data.append(PostBreachData(self.name, self.command, results))
+        return self.pba_data
 
-    def modify_shell_startup_PBA_list(self):
+    @classmethod
-        return self.ShellStartupPBAGenerator().get_modify_shell_startup_pbas()
+    def modify_shell_startup_PBA_list(cls):
+        return cls.ShellStartupPBAGenerator.get_modify_shell_startup_pbas()
 
     class ShellStartupPBAGenerator:
-        def get_modify_shell_startup_pbas(self):
+        @classmethod
+        def get_modify_shell_startup_pbas(cls):
             (cmds_for_linux, shell_startup_files_for_linux, usernames_for_linux), (
                 cmds_for_windows,
                 shell_startup_files_per_user_for_windows,
@@ -43,14 +48,14 @@ class ModifyShellStartupFiles(PBA):
 
             for startup_file_per_user in shell_startup_files_per_user_for_windows:
                 windows_cmds = " ".join(cmds_for_windows).format(startup_file_per_user)
-                pbas.append(self.ModifyShellStartupFile(linux_cmds="", windows_cmds=windows_cmds))
+                pbas.append(cls.ModifyShellStartupFile(linux_cmds="", windows_cmds=windows_cmds))
 
             for username in usernames_for_linux:
                 for shell_startup_file in shell_startup_files_for_linux:
                     linux_cmds = (
                         " ".join(cmds_for_linux).format(shell_startup_file).format(username)
                     )
-                    pbas.append(self.ModifyShellStartupFile(linux_cmds=linux_cmds, windows_cmds=""))
+                    pbas.append(cls.ModifyShellStartupFile(linux_cmds=linux_cmds, windows_cmds=""))
 
             return pbas
 

monkey/infection_monkey/post_breach/actions/schedule_jobs.py

@@ -23,3 +23,4 @@ class ScheduleJobs(PBA):
     def run(self):
         super(ScheduleJobs, self).run()
         remove_scheduled_jobs()
+        return self.pba_data

monkey/infection_monkey/post_breach/actions/use_signed_scripts.py

@@ -18,13 +18,14 @@ class SignedScriptProxyExecution(PBA):
         super().__init__(POST_BREACH_SIGNED_SCRIPT_PROXY_EXEC, windows_cmd=" ".join(windows_cmds))
 
     def run(self):
+        original_comspec = ""
         try:
-            original_comspec = ""
             if is_windows_os():
                 original_comspec = subprocess.check_output(  # noqa: DUO116
                     "if defined COMSPEC echo %COMSPEC%", shell=True
                 ).decode()
             super().run()
+            return self.pba_data
         except Exception as e:
             logger.warning(
                 f"An exception occurred on running PBA "

monkey/infection_monkey/post_breach/pba.py

@@ -1,9 +1,10 @@
 import logging
 import subprocess
+from typing import Iterable
 
 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.post_breach_telem import PostBreachTelem
 from infection_monkey.utils.environment import is_windows_os
 
 logger = logging.getLogger(__name__)
@@ -23,8 +24,9 @@ class PBA:
         """
         self.command = PBA.choose_command(linux_cmd, windows_cmd)
         self.name = name
+        self.pba_data = []
 
-    def run(self):
+    def run(self) -> Iterable[PostBreachData]:
         """
         Runs post breach action command
         """
@@ -35,7 +37,8 @@ class PBA:
                 T1064Telem(
                     ScanStatus.USED, f"Scripts were used to execute {self.name} post breach action."
                 ).send()
-            PostBreachTelem(self, result).send()
+            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.")
 

monkey/infection_monkey/puppet/mock_puppet.py

@@ -53,9 +53,9 @@ class MockPuppet(IPuppet):
         logger.debug(f"run_pba({name}, {options})")
 
         if name == "AccountDiscovery":
-            return PostBreachData(name, "pba command 1", ["pba result 1", True])
+            yield PostBreachData(name, "pba command 1", ["pba result 1", True])
         else:
-            return PostBreachData(name, "pba command 2", ["pba result 2", False])
+            yield PostBreachData(name, "pba command 2", ["pba result 2", False])
 
     def ping(self, host: str, timeout: float = 1) -> PingScanData:
         logger.debug(f"run_ping({host}, {timeout})")

vulture_allowlist.py

@@ -92,6 +92,7 @@ AccountDiscovery  # unused class (monkey/infection_monkey/post_breach/actions/di
 ModifyShellStartupFiles  # unused class (monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py:11)
 Timestomping  # unused class (monkey/infection_monkey/post_breach/actions/timestomping.py:6)
 SignedScriptProxyExecution  # unused class (monkey/infection_monkey/post_breach/actions/use_signed_scripts.py:15)
+ProcessListCollection  # unused class (monkey/infection_monkey/post_breach/actions/collect_processes_list.py:19)
 EnvironmentCollector  # unused class (monkey/infection_monkey/system_info/collectors/environment_collector.py:19)
 HostnameCollector  # unused class (monkey/infection_monkey/system_info/collectors/hostname_collector.py:10)
 _.representations  # unused attribute (monkey/monkey_island/cc/app.py:180)