Code design changes

monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py

@@ -1,6 +1,6 @@
 from common.data.post_breach_consts import POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION
 from infection_monkey.post_breach.pba import PBA
-from infection_monkey.utils.shell_startup_files_modification import\
+from infection_monkey.post_breach.shell_startup_files.shell_startup_files_modification import\
     get_commands_to_modify_shell_startup_files
 from infection_monkey.utils.environment import is_windows_os
 
@@ -12,21 +12,31 @@ class ModifyShellStartupFiles(PBA):
     and profile.ps1 in windows.
     """
 
-    def __init__(self):
-        super(ModifyShellStartupFiles, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION)
-
     def run(self):
+        [pba.run() for pba in self.modify_shell_startup_PBA_list()]
+
+    def modify_shell_startup_PBA_list(self):
+        return ShellStartupPBAGenerator.get_modify_shell_startup_pbas()
+
+
+class ShellStartupPBAGenerator():
+    def get_modify_shell_startup_pbas():
         (cmds_for_linux, shell_startup_files_for_linux), windows_cmds = get_commands_to_modify_shell_startup_files()
 
-        if is_windows_os():
+        pbas = [ModifyShellStartupFile(linux_cmds='', windows_cmds=windows_cmds)]
-            super(ModifyShellStartupFiles, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION,
+
-                                                          linux_cmd='',  # windows so won't matter
-                                                          windows_cmd=windows_cmds)
-            super(ModifyShellStartupFiles, self).run()
-        else:
         for shell_startup_file in shell_startup_files_for_linux:
             linux_cmds = ' '.join(cmds_for_linux).format(shell_startup_file)
-                super(ModifyShellStartupFiles, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION,
+            pbas.append(ModifyShellStartupFile(linux_cmds=linux_cmds, windows_cmds=''))
+
+        return pbas
+
+
+class ModifyShellStartupFile(PBA):
+    def __init__(self, linux_cmds, windows_cmds):
+        super(ModifyShellStartupFile, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION,
                                                      linux_cmd=linux_cmds,
                                                      windows_cmd=windows_cmds)
-                super(ModifyShellStartupFiles, self).run()
+
+    def run(self):
+        super(ModifyShellStartupFile, self).run()

monkey/infection_monkey/post_breach/pba.py

@@ -57,6 +57,7 @@ class PBA(Plugin):
         """
         Runs post breach action command
         """
+        if self.command:
             exec_funct = self._execute_default
             result = exec_funct()
             if self.scripts_were_used_successfully(result):