From f6556704d667c7f26764adfa0ea596280734ae98 Mon Sep 17 00:00:00 2001 From: Shreya Date: Wed, 22 Jul 2020 01:57:11 +0530 Subject: [PATCH] Modify `run()` for "modify shell startup files" PBA --- .../actions/modify_shell_startup_files.py | 13 +------------ monkey/infection_monkey/post_breach/pba.py | 7 +++++-- 2 files changed, 6 insertions(+), 14 deletions(-) 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 3973724e0..352610c6b 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 @@ -21,7 +21,7 @@ class ModifyShellStartupFiles(PBA): super().__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION) def run(self): - results = [pba.run() for pba in self.modify_shell_startup_PBA_list()] + results = [pba.run(return_result=True) for pba in self.modify_shell_startup_PBA_list()] PostBreachTelem(self, results).send() def modify_shell_startup_PBA_list(self): @@ -51,14 +51,3 @@ class ModifyShellStartupFiles(PBA): super().__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION, linux_cmd=linux_cmds, windows_cmd=windows_cmds) - - def run(self): - if self.command: - try: - output = subprocess.check_output(self.command, stderr=subprocess.STDOUT, shell=True).decode() - if not output: - output = EXECUTION_WITHOUT_OUTPUT - return output, True - except subprocess.CalledProcessError as e: - # Return error output of the command - return e.output.decode(), False diff --git a/monkey/infection_monkey/post_breach/pba.py b/monkey/infection_monkey/post_breach/pba.py index a6a89edf8..0660449ad 100644 --- a/monkey/infection_monkey/post_breach/pba.py +++ b/monkey/infection_monkey/post_breach/pba.py @@ -54,7 +54,7 @@ class PBA(Plugin): """ return class_name in WormConfiguration.post_breach_actions - def run(self): + def run(self, return_result=False): """ Runs post breach action command """ @@ -63,7 +63,10 @@ class PBA(Plugin): result = exec_funct() if self.scripts_were_used_successfully(result): T1064Telem(ScanStatus.USED, f"Scripts were used to execute {self.name} post breach action.").send() - PostBreachTelem(self, result).send() + if return_result: + return result + else: + PostBreachTelem(self, result).send() else: LOG.debug(f"No command available for PBA '{self.name}' on current OS, skipping.")