CR changes

- Added nested classes
- Extracted repetitive code

monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py

@@ -18,48 +18,47 @@ class ModifyShellStartupFiles(PBA):
     """
 
     def __init__(self):
-        super(ModifyShellStartupFiles, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION)
+        super().__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION)
 
     def run(self):
         results = [pba.run() for pba in self.modify_shell_startup_PBA_list()]
         PostBreachTelem(self, results).send()
 
     def modify_shell_startup_PBA_list(self):
-        return ShellStartupPBAGenerator.get_modify_shell_startup_pbas()
+        return self.ShellStartupPBAGenerator().get_modify_shell_startup_pbas()
 
+    class ShellStartupPBAGenerator():
+        def get_modify_shell_startup_pbas(self):
+            (cmds_for_linux, shell_startup_files_for_linux, usernames_for_linux),\
+             (cmds_for_windows, shell_startup_files_per_user_for_windows) =\
+             get_commands_to_modify_shell_startup_files()
 
-class ShellStartupPBAGenerator():
+            pbas = []
-    def get_modify_shell_startup_pbas():
-        (cmds_for_linux, shell_startup_files_for_linux, usernames_for_linux),\
-         (cmds_for_windows, shell_startup_files_per_user_for_windows) = get_commands_to_modify_shell_startup_files()
 
-        pbas = []
+            for startup_file_per_user in shell_startup_files_per_user_for_windows:
+                windows_cmds = ' '.join(cmds_for_windows).format(startup_file_per_user)
+                pbas.append(self.ModifyShellStartupFile(linux_cmds='', windows_cmds=['powershell.exe', windows_cmds]))
 
-        for startup_file_per_user in shell_startup_files_per_user_for_windows:
+            for username in usernames_for_linux:
-            windows_cmds = ' '.join(cmds_for_windows).format(startup_file_per_user)
+                for shell_startup_file in shell_startup_files_for_linux:
-            pbas.append(ModifyShellStartupFile(linux_cmds='', windows_cmds=['powershell.exe', windows_cmds]))
+                    linux_cmds = ' '.join(cmds_for_linux).format(shell_startup_file).format(username)
+                    pbas.append(self.ModifyShellStartupFile(linux_cmds=linux_cmds, windows_cmds=''))
 
-        for username in usernames_for_linux:
+            return pbas
-            for shell_startup_file in shell_startup_files_for_linux:
-                linux_cmds = ' '.join(cmds_for_linux).format(shell_startup_file).format(username)
-                pbas.append(ModifyShellStartupFile(linux_cmds=linux_cmds, windows_cmds=''))
 
-        return pbas
+        class ModifyShellStartupFile(PBA):
+            def __init__(self, linux_cmds, windows_cmds):
+                super().__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION,
+                                 linux_cmd=linux_cmds,
+                                 windows_cmd=windows_cmds)
 
-
+            def run(self):
-class ModifyShellStartupFile(PBA):
+                if self.command:
-    def __init__(self, linux_cmds, windows_cmds):
+                    try:
-        super(ModifyShellStartupFile, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION,
+                        output = subprocess.check_output(self.command, stderr=subprocess.STDOUT, shell=True).decode()
-                                                     linux_cmd=linux_cmds,
+                        if not output:
-                                                     windows_cmd=windows_cmds)
+                            output = EXECUTION_WITHOUT_OUTPUT
-
+                        return output, True
-    def run(self):
+                    except subprocess.CalledProcessError as e:
-        if self.command:
+                        # Return error output of the command
-            try:
+                        return e.output.decode(), False
-                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

monkey/monkey_island/cc/services/attack/technique_reports/T1156.py

@@ -1,8 +1,9 @@
 from common.data.post_breach_consts import \
     POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION
-from common.utils.attack_utils import ScanStatus
 from monkey_island.cc.database import mongo
 from monkey_island.cc.services.attack.technique_reports import AttackTechnique
+from monkey_island.cc.services.attack.technique_reports.technique_report_tools import \
+    extract_shell_startup_files_modification_info, get_shell_startup_files_modification_status
 
 __author__ = "shreyamalviya"
 
@@ -26,20 +27,10 @@ class T1156(AttackTechnique):
 
         shell_startup_files_modification_info = list(mongo.db.telemetry.aggregate(T1156.query))
 
-        bash_startup_modification_info = []
+        bash_startup_modification_info =\
-        for shell_startup_file_result in shell_startup_files_modification_info[0]['result']:
+            extract_shell_startup_files_modification_info(shell_startup_files_modification_info, [".bash", ".profile"])
-            # only want bash startup files
-            if any(file_name in shell_startup_file_result[0] for file_name in [".bash", ".profile"]):
-                bash_startup_modification_info.append({
-                    'machine': shell_startup_files_modification_info[0]['machine'],
-                    'result': shell_startup_file_result
-                    })
 
-        status = []
+        status = get_shell_startup_files_modification_status(bash_startup_modification_info)
-        for bash_startup_file in bash_startup_modification_info:
-            status.append(bash_startup_file['result'][1])
-        status = (ScanStatus.USED.value if any(status) else ScanStatus.SCANNED.value)\
-            if status else ScanStatus.UNSCANNED.value
 
         data.update(T1156.get_base_data_by_status(status))
         data.update({'info': bash_startup_modification_info})

monkey/monkey_island/cc/services/attack/technique_reports/T1504.py

@@ -1,8 +1,9 @@
 from common.data.post_breach_consts import \
     POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION
-from common.utils.attack_utils import ScanStatus
 from monkey_island.cc.database import mongo
 from monkey_island.cc.services.attack.technique_reports import AttackTechnique
+from monkey_island.cc.services.attack.technique_reports.technique_report_tools import \
+    extract_shell_startup_files_modification_info, get_shell_startup_files_modification_status
 
 __author__ = "shreyamalviya"
 
@@ -26,20 +27,10 @@ class T1504(AttackTechnique):
 
         shell_startup_files_modification_info = list(mongo.db.telemetry.aggregate(T1504.query))
 
-        powershell_startup_modification_info = []
+        powershell_startup_modification_info =\
-        for shell_startup_file_result in shell_startup_files_modification_info[0]['result']:
+            extract_shell_startup_files_modification_info(shell_startup_files_modification_info, ["profile.ps1"])
-            # only want powershell startup files
-            if "profile.ps1" in shell_startup_file_result[0]:
-                powershell_startup_modification_info.append({
-                    'machine': shell_startup_files_modification_info[0]['machine'],
-                    'result': shell_startup_file_result
-                    })
 
-        status = []
+        status = get_shell_startup_files_modification_status(powershell_startup_modification_info)
-        for powershell_startup_file in powershell_startup_modification_info:
-            status.append(powershell_startup_file['result'][1])
-        status = (ScanStatus.USED.value if any(status) else ScanStatus.SCANNED.value)\
-            if status else ScanStatus.UNSCANNED.value
 
         data.update(T1504.get_base_data_by_status(status))
         data.update({'info': powershell_startup_modification_info})

monkey/monkey_island/cc/services/attack/technique_reports/technique_report_tools.py

@@ -1,4 +1,5 @@
 from monkey_island.cc.encryptor import encryptor
+from common.utils.attack_utils import ScanStatus
 
 
 def parse_creds(attempt):
@@ -44,3 +45,23 @@ def censor_hash(hash_, plain_chars=5):
         return ""
     hash_ = encryptor.dec(hash_)
     return hash_[0: plain_chars] + ' ...'
+
+
+def extract_shell_startup_files_modification_info(shell_startup_files_modification_info, required_file_names):
+    required_shell_startup_files_modification_info = []
+    for shell_startup_file_result in shell_startup_files_modification_info[0]['result']:
+        if any(file_name in shell_startup_file_result[0] for file_name in required_file_names):
+            shell_startup_files_modification_info.append({
+                'machine': shell_startup_files_modification_info[0]['machine'],
+                'result': shell_startup_file_result
+                })
+    return required_shell_startup_files_modification_info
+
+
+def get_shell_startup_files_modification_status(shell_startup_files_modification_info):
+    status = []
+    for startup_file in shell_startup_files_modification_info:
+        status.append(startup_file['result'][1])
+    status = (ScanStatus.USED.value if any(status) else ScanStatus.SCANNED.value)\
+        if status else ScanStatus.UNSCANNED.value
+    return status