From 5a7e8a0b08d33faa87c732064ea00cb251e2f72c Mon Sep 17 00:00:00 2001
From: Shreya <shreya.malviya@gmail.com>
Date: Fri, 24 Jul 2020 13:13:11 +0530
Subject: [PATCH] Override `run()` for PBA "modify shell startup files" +
 assign EXECUTION_WITHOUT_OUTPUT to all relevant PBAs on the island side

---
 .../post_breach/actions/modify_shell_startup_files.py |  9 +++++++++
 monkey/infection_monkey/post_breach/pba.py            | 11 ++---------
 .../cc/services/telemetry/processing/post_breach.py   |  8 ++++++++
 3 files changed, 19 insertions(+), 9 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 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