CR changes

- Added nested classes
- Extracted repetitive code

monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py

@@ -18,38 +18,37 @@ 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():
+    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()
+             (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(ModifyShellStartupFile(linux_cmds='', windows_cmds=['powershell.exe', windows_cmds]))
+                pbas.append(self.ModifyShellStartupFile(linux_cmds='', windows_cmds=['powershell.exe', windows_cmds]))
 
             for username in usernames_for_linux:
                 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=''))
+                    pbas.append(self.ModifyShellStartupFile(linux_cmds=linux_cmds, windows_cmds=''))
 
             return pbas
 
-
-class ModifyShellStartupFile(PBA):
+        class ModifyShellStartupFile(PBA):
             def __init__(self, linux_cmds, windows_cmds):
-        super(ModifyShellStartupFile, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION,
+                super().__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION,
                                  linux_cmd=linux_cmds,
                                  windows_cmd=windows_cmds)
 

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 = []
-        for shell_startup_file_result in shell_startup_files_modification_info[0]['result']:
-            # 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
-                    })
+        bash_startup_modification_info =\
+            extract_shell_startup_files_modification_info(shell_startup_files_modification_info, [".bash", ".profile"])
 
-        status = []
-        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
+        status = get_shell_startup_files_modification_status(bash_startup_modification_info)
 
         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 = []
-        for shell_startup_file_result in shell_startup_files_modification_info[0]['result']:
-            # 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
-                    })
+        powershell_startup_modification_info =\
+            extract_shell_startup_files_modification_info(shell_startup_files_modification_info, ["profile.ps1"])
 
-        status = []
-        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
+        status = get_shell_startup_files_modification_status(powershell_startup_modification_info)
 
         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