Agent: Modify PBAs to return Iterable[PostBreachData]

monkey/infection_monkey/post_breach/actions/clear_command_history.py

@@ -15,8 +15,9 @@ class ClearCommandHistory(PBA):
     def run(self):
         results = [pba.run() for pba in self.clear_command_history_PBA_list()]
         if results:
-            # Note: `self.command` is empty here
-            yield PostBreachData(self.name, self.command, 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):
         return self.CommandHistoryPBAGenerator().get_clear_command_history_pbas()

monkey/infection_monkey/post_breach/actions/collect_processes_list.py

@@ -53,4 +53,5 @@ class ProcessListCollection(PBA):
                 continue
 
         # No command here; used psutil
-        yield PostBreachData(self.name, "", (processes, success_state))
+        self.pba_data.append(PostBreachData(self.name, "", (processes, success_state)))
+        return self.pba_data

monkey/infection_monkey/post_breach/actions/communicate_as_backdoor_user.py

@@ -54,11 +54,13 @@ class CommunicateAsBackdoorUser(PBA):
                 )
                 # `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
-                yield PostBreachData(self.name, "", result)
+                self.pba_data.append(PostBreachData(self.name, "", result))
         except subprocess.CalledProcessError as e:
-            yield PostBreachData(self.name, "", (e.output.decode(), False))
+            self.pba_data.append(PostBreachData(self.name, "", (e.output.decode(), False)))
         except NewUserError as e:
-            yield PostBreachData(self.name, "", (str(e), False))
+            self.pba_data.append(PostBreachData(self.name, "", (str(e), False)))
+        finally:
+            return self.pba_data
 
     @staticmethod
     def get_random_new_user_name():

monkey/infection_monkey/post_breach/actions/hide_files.py

@@ -29,12 +29,14 @@ class HiddenFiles(PBA):
                 linux_cmd=" ".join(linux_cmds),
                 windows_cmd=windows_cmds,
             )
-            yield super(HiddenFiles, self).run()
+            super(HiddenFiles, self).run()
 
         if is_windows_os():  # use winAPI
             result, status = get_winAPI_to_hide_files()
             # no command here, used WinAPI
-            yield PostBreachData(self.name, "", (result, status))
+            self.pba_data.append(PostBreachData(self.name, "", (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

@@ -29,7 +29,8 @@ class ModifyShellStartupFiles(PBA):
             ]
         # `command` is empty here since multiple commands were run and the results
         # were aggregated to send the telemetry just once
-        yield PostBreachData(self.name, "", results).send()
+        self.pba_data.append(PostBreachData(self.name, "", results))
+        return self.pba_data
 
     def modify_shell_startup_PBA_list(self):
         return self.ShellStartupPBAGenerator().get_modify_shell_startup_pbas()

monkey/infection_monkey/post_breach/actions/schedule_jobs.py

@@ -21,6 +21,6 @@ class ScheduleJobs(PBA):
         )
 
     def run(self):
-        post_breach_data = super(ScheduleJobs, self).run()
+        super(ScheduleJobs, self).run()
         remove_scheduled_jobs()
-        yield post_breach_data
+        return self.pba_data

monkey/infection_monkey/post_breach/actions/use_signed_scripts.py

@@ -24,7 +24,8 @@ class SignedScriptProxyExecution(PBA):
                 original_comspec = subprocess.check_output(  # noqa: DUO116
                     "if defined COMSPEC echo %COMSPEC%", shell=True
                 ).decode()
-            yield super().run()
+            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,5 +1,6 @@
 import logging
 import subprocess
+from typing import Iterable
 
 from common.utils.attack_utils import ScanStatus
 from infection_monkey.i_puppet.i_puppet import PostBreachData
@@ -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()
-            yield PostBreachData(self.name, self.command, result)
+            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.")