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 9baa3dc67..036c32d25 100644 --- a/monkey/infection_monkey/post_breach/actions/clear_command_history.py +++ b/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() 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 782c771dc..5f18e0e33 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,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 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 36c96b126..73ef0fa3b 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,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(): diff --git a/monkey/infection_monkey/post_breach/actions/hide_files.py b/monkey/infection_monkey/post_breach/actions/hide_files.py index 6bbeefa68..1a2f3472d 100644 --- a/monkey/infection_monkey/post_breach/actions/hide_files.py +++ b/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 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 75b2e1a55..bb1a653f8 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,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() diff --git a/monkey/infection_monkey/post_breach/actions/schedule_jobs.py b/monkey/infection_monkey/post_breach/actions/schedule_jobs.py index 8846efcf9..37649488b 100644 --- a/monkey/infection_monkey/post_breach/actions/schedule_jobs.py +++ b/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 diff --git a/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py b/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py index 984fc0f66..f6066fecb 100644 --- a/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py +++ b/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 " diff --git a/monkey/infection_monkey/post_breach/pba.py b/monkey/infection_monkey/post_breach/pba.py index 449c06186..8b50f08ba 100644 --- a/monkey/infection_monkey/post_breach/pba.py +++ b/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.")