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 85f039628..fdc7e18dc 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 @@ -47,3 +47,12 @@ 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() + return output, True + except subprocess.CalledProcessError as e: + # Return error output of the command + return e.output.decode(), False \ No newline at end of file diff --git a/monkey/infection_monkey/post_breach/pba.py b/monkey/infection_monkey/post_breach/pba.py index 0660449ad..93d10d45e 100644 --- a/monkey/infection_monkey/post_breach/pba.py +++ b/monkey/infection_monkey/post_breach/pba.py @@ -13,8 +13,6 @@ LOG = logging.getLogger(__name__) __author__ = 'VakarisZ' -EXECUTION_WITHOUT_OUTPUT = "(PBA execution produced no output)" - class PBA(Plugin): """ @@ -54,7 +52,7 @@ class PBA(Plugin): """ return class_name in WormConfiguration.post_breach_actions - def run(self, return_result=False): + def run(self): """ Runs post breach action command """ @@ -63,10 +61,7 @@ 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() - if return_result: - return result - else: - PostBreachTelem(self, result).send() + PostBreachTelem(self, result).send() else: LOG.debug(f"No command available for PBA '{self.name}' on current OS, skipping.") @@ -93,8 +88,6 @@ class PBA(Plugin): """ 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 diff --git a/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py b/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py index 63cdfff48..446ac0372 100644 --- a/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py +++ b/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py @@ -8,6 +8,8 @@ from monkey_island.cc.models import Monkey from monkey_island.cc.services.telemetry.zero_trust_tests.communicate_as_new_user import \ test_new_user_communication +EXECUTION_WITHOUT_OUTPUT = "(PBA execution produced no output)" + def process_communicate_as_new_user_telemetry(telemetry_json): current_monkey = Monkey.get_single_monkey_by_guid(telemetry_json['monkey_guid']) @@ -38,10 +40,16 @@ def process_post_breach_telemetry(telemetry_json): if type(telemetry_json['data']) is list: for pba_data in telemetry_json['data']: + modify_blank_outputs(pba_data) mongo.db.monkey.update( {'guid': telemetry_json['monkey_guid']}, {'$push': {'pba_results': pba_data}}) else: + modify_blank_outputs(telemetry_json['data']) mongo.db.monkey.update( {'guid': telemetry_json['monkey_guid']}, {'$push': {'pba_results': telemetry_json['data']}}) + + def modify_blank_outputs(data): + if not data['result']: + data['result'] = EXECUTION_WITHOUT_OUTPUT