From 0c60ad16aed66d1b2fae83b6a0cfb81fe09a7094 Mon Sep 17 00:00:00 2001
From: Shreya <shreya.malviya@gmail.com>
Date: Wed, 17 Jun 2020 11:49:03 +0530
Subject: [PATCH] Code design changes

---
 .../actions/modify_shell_startup_files.py     | 42 ++++++++++++-------
 monkey/infection_monkey/post_breach/pba.py    | 11 ++---
 2 files changed, 32 insertions(+), 21 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 a526acad6..8417849db 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
@@ -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():
-            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,
-                                                              linux_cmd=linux_cmds,
-                                                              windows_cmd=windows_cmds)
-                super(ModifyShellStartupFiles, self).run()
+        pbas = [ModifyShellStartupFile(linux_cmds='', windows_cmds=windows_cmds)]
+
+        for shell_startup_file in shell_startup_files_for_linux:
+            linux_cmds = ' '.join(cmds_for_linux).format(shell_startup_file)
+            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)
+
+    def run(self):
+        super(ModifyShellStartupFile, self).run()
diff --git a/monkey/infection_monkey/post_breach/pba.py b/monkey/infection_monkey/post_breach/pba.py
index 3d8da9dab..e5a91b733 100644
--- a/monkey/infection_monkey/post_breach/pba.py
+++ b/monkey/infection_monkey/post_breach/pba.py
@@ -57,11 +57,12 @@ class PBA(Plugin):
         """
         Runs post breach action command
         """
-        exec_funct = self._execute_default
-        result = exec_funct()
-        if self.scripts_were_used_successfully(result):
-            T1064Telem(ScanStatus.USED, "Scripts were used to execute %s post breach action." % self.name).send()
-        PostBreachTelem(self, result).send()
+        if self.command:
+            exec_funct = self._execute_default
+            result = exec_funct()
+            if self.scripts_were_used_successfully(result):
+                T1064Telem(ScanStatus.USED, "Scripts were used to execute %s post breach action." % self.name).send()
+            PostBreachTelem(self, result).send()
 
     def is_script(self):
         """