diff --git a/monkey/infection_monkey/post_breach/actions/clear_command_history.py b/monkey/infection_monkey/post_breach/actions/clear_command_history.py index 03cb77be0..9baa3dc67 100644 --- a/monkey/infection_monkey/post_breach/actions/clear_command_history.py +++ b/monkey/infection_monkey/post_breach/actions/clear_command_history.py @@ -16,7 +16,7 @@ class ClearCommandHistory(PBA): results = [pba.run() for pba in self.clear_command_history_PBA_list()] if results: # Note: `self.command` is empty here - return PostBreachData(self.name, self.command, results) + yield PostBreachData(self.name, self.command, results) def clear_command_history_PBA_list(self): return self.CommandHistoryPBAGenerator().get_clear_command_history_pbas() diff --git a/monkey/infection_monkey/post_breach/actions/collect_processes_list.py b/monkey/infection_monkey/post_breach/actions/collect_processes_list.py index 260d4bf18..782c771dc 100644 --- a/monkey/infection_monkey/post_breach/actions/collect_processes_list.py +++ b/monkey/infection_monkey/post_breach/actions/collect_processes_list.py @@ -53,4 +53,4 @@ class ProcessListCollection(PBA): continue # No command here; used psutil - return PostBreachData(self.name, "", (processes, success_state)) + yield PostBreachData(self.name, "", (processes, success_state)) diff --git a/monkey/infection_monkey/post_breach/actions/communicate_as_backdoor_user.py b/monkey/infection_monkey/post_breach/actions/communicate_as_backdoor_user.py index d93be17e1..36c96b126 100644 --- a/monkey/infection_monkey/post_breach/actions/communicate_as_backdoor_user.py +++ b/monkey/infection_monkey/post_breach/actions/communicate_as_backdoor_user.py @@ -54,11 +54,11 @@ 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 - return PostBreachData(self.name, "", result) + yield PostBreachData(self.name, "", result) except subprocess.CalledProcessError as e: - return PostBreachData(self.name, "", (e.output.decode(), False)) + yield PostBreachData(self.name, "", (e.output.decode(), False)) except NewUserError as e: - return PostBreachData(self.name, "", (str(e), False)) + yield PostBreachData(self.name, "", (str(e), False)) @staticmethod def get_random_new_user_name(): diff --git a/monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py b/monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py index ebaf9dfc1..75b2e1a55 100644 --- a/monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py +++ b/monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py @@ -29,7 +29,7 @@ class ModifyShellStartupFiles(PBA): ] # `command` is empty here since multiple commands were run and the results # were aggregated to send the telemetry just once - return PostBreachData(self.name, "", results).send() + yield PostBreachData(self.name, "", results).send() def modify_shell_startup_PBA_list(self): return self.ShellStartupPBAGenerator().get_modify_shell_startup_pbas() diff --git a/monkey/infection_monkey/post_breach/pba.py b/monkey/infection_monkey/post_breach/pba.py index ab3a004f0..449c06186 100644 --- a/monkey/infection_monkey/post_breach/pba.py +++ b/monkey/infection_monkey/post_breach/pba.py @@ -35,7 +35,7 @@ class PBA: T1064Telem( ScanStatus.USED, f"Scripts were used to execute {self.name} post breach action." ).send() - return PostBreachData(self.name, self.command, result) + yield PostBreachData(self.name, self.command, result) else: logger.debug(f"No command available for PBA '{self.name}' on current OS, skipping.") diff --git a/monkey/infection_monkey/puppet/mock_puppet.py b/monkey/infection_monkey/puppet/mock_puppet.py index 0196076ad..5f707acd7 100644 --- a/monkey/infection_monkey/puppet/mock_puppet.py +++ b/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})")