From ec2b2beca5b15aee5a91b77b27431dc94bdaddf7 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 28 Mar 2022 17:16:50 +0530 Subject: [PATCH] Agent: Modify PBAs to yield PostBreachData instead of returning it This is done mainly because of the hide files PBA which needs to send telemetry two times. It also makes more sense to do it this way so that it's easier to send telemetry multiple times in any PBA. --- .../post_breach/actions/clear_command_history.py | 2 +- .../post_breach/actions/collect_processes_list.py | 2 +- .../post_breach/actions/communicate_as_backdoor_user.py | 6 +++--- .../post_breach/actions/modify_shell_startup_files.py | 2 +- monkey/infection_monkey/post_breach/pba.py | 2 +- monkey/infection_monkey/puppet/mock_puppet.py | 4 ++-- 6 files changed, 9 insertions(+), 9 deletions(-) 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})")