Refactored to automatically check if post breach should run or not

monkey/infection_monkey/post_breach/actions/add_user.py

@@ -14,12 +14,6 @@ WINDOWS_COMMANDS = ['net', 'user', WormConfiguration.user_to_add,
                     '/add', '/ACTIVE:NO']
 
 
-class BackdoorUser(object):
+class BackdoorUser(PBA):
-    PBA_NAME = "Backdoor user"
-
     def __init__(self):
-        pass
+        super(BackdoorUser, self).__init__("Backdoor user", linux_cmd=LINUX_COMMANDS, windows_cmd=WINDOWS_COMMANDS)
-
-    @staticmethod
-    def get_pba():
-        return PBA.default_get_pba(BackdoorUser.PBA_NAME, BackdoorUser, LINUX_COMMANDS, WINDOWS_COMMANDS)

monkey/infection_monkey/post_breach/actions/users_custom_pba.py

@@ -1,4 +1,3 @@
-import requests
 import os
 import logging
 
@@ -24,14 +23,44 @@ class UsersPBA(PBA):
     """
     Defines user's configured post breach action.
     """
-    def __init__(self, command="", filename=""):
+    def __init__(self):
-        self.filename = filename
+        super(UsersPBA, self).__init__("File execution")
-        super(UsersPBA, self).__init__("File execution", command)
+        self.filename = ''
+        if not is_windows_os():
+            # Add linux commands to PBA's
+            if WormConfiguration.PBA_linux_filename:
+                if WormConfiguration.custom_PBA_linux_cmd:
+                    # Add change dir command, because user will try to access his file
+                    self.command = (DIR_CHANGE_LINUX % get_monkey_dir_path()) + WormConfiguration.custom_PBA_linux_cmd
+                    self.filename = WormConfiguration.PBA_linux_filename
+                else:
+                    file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_linux_filename)
+                    self.command = DEFAULT_LINUX_COMMAND.format(file_path)
+                    self.filename = WormConfiguration.PBA_linux_filename
+            elif WormConfiguration.custom_PBA_linux_cmd:
+                self.command = WormConfiguration.custom_PBA_linux_cmd
+        else:
+            # Add windows commands to PBA's
+            if WormConfiguration.PBA_windows_filename:
+                if WormConfiguration.custom_PBA_windows_cmd:
+                    # Add change dir command, because user will try to access his file
+                    self.command = (DIR_CHANGE_WINDOWS % get_monkey_dir_path()) + WormConfiguration.custom_PBA_windows_cmd
+                    self.filename = WormConfiguration.PBA_windows_filename
+                else:
+                    file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_windows_filename)
+                    self.command = DEFAULT_WINDOWS_COMMAND.format(file_path)
+                    self.filename = WormConfiguration.PBA_windows_filename
+            elif WormConfiguration.custom_PBA_windows_cmd:
+                self.command = WormConfiguration.custom_PBA_windows_cmd
 
     def _execute_default(self):
+        if self.filename:
             UsersPBA.download_pba_file(get_monkey_dir_path(), self.filename)
         return super(UsersPBA, self)._execute_default()
 
+    def should_run(self, class_name):
+        return self.command
+
     @staticmethod
     def download_pba_file(dst_dir, filename):
         """
@@ -53,38 +82,3 @@ class UsersPBA(PBA):
         except IOError as e:
             LOG.error("Can not upload post breach file to target machine: %s" % e)
             return False
-
-    @staticmethod
-    def get_pba():
-        """
-        Creates post breach actions depending on users input into 'custom post breach' config section
-        :return: List of PBA objects ([user's file execution PBA, user's command execution PBA])
-        """
-        command_pba_name = "Custom command"
-
-        if not is_windows_os():
-            # Add linux commands to PBA's
-            if WormConfiguration.PBA_linux_filename:
-                if WormConfiguration.custom_PBA_linux_cmd:
-                    # Add change dir command, because user will try to access his file
-                    command = (DIR_CHANGE_LINUX % get_monkey_dir_path()) + WormConfiguration.custom_PBA_linux_cmd
-                    return UsersPBA(command, WormConfiguration.PBA_linux_filename)
-                else:
-                    file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_linux_filename)
-                    command = DEFAULT_LINUX_COMMAND.format(file_path)
-                    return UsersPBA(command, WormConfiguration.PBA_linux_filename)
-            elif WormConfiguration.custom_PBA_linux_cmd:
-                return PBA(name=command_pba_name, command=WormConfiguration.custom_PBA_linux_cmd)
-        else:
-            # Add windows commands to PBA's
-            if WormConfiguration.PBA_windows_filename:
-                if WormConfiguration.custom_PBA_windows_cmd:
-                    # Add change dir command, because user will try to access his file
-                    command = (DIR_CHANGE_WINDOWS % get_monkey_dir_path()) + WormConfiguration.custom_PBA_windows_cmd
-                    return UsersPBA(command, WormConfiguration.PBA_windows_filename)
-                else:
-                    file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_windows_filename)
-                    command = DEFAULT_WINDOWS_COMMAND.format(file_path)
-                    return UsersPBA(command, WormConfiguration.PBA_windows_filename)
-            elif WormConfiguration.custom_PBA_windows_cmd:
-                return PBA(name=command_pba_name, command=WormConfiguration.custom_PBA_windows_cmd)

monkey/infection_monkey/post_breach/pba.py

@@ -2,7 +2,7 @@ import logging
 import subprocess
 from infection_monkey.control import ControlClient
 from infection_monkey.utils import is_windows_os
-from infection_monkey.config import WormConfiguration, GUID
+from infection_monkey.config import WormConfiguration
 
 
 LOG = logging.getLogger(__name__)
@@ -14,38 +14,28 @@ class PBA(object):
     """
     Post breach action object. Can be extended to support more than command execution on target machine.
     """
-    def __init__(self, name="unknown", command=""):
+    def __init__(self, name="unknown", linux_cmd="", windows_cmd=""):
         """
         :param name: Name of post breach action.
         :param command: Command that will be executed on breached machine
         """
-        self.command = command
+        self.command = PBA.choose_command(linux_cmd, windows_cmd)
         self.name = name
 
-    @staticmethod
+    def get_pba(self):
-    def get_pba():
         """
-        Should be overridden by all child classes.
+        This method returns a PBA object based on a worm's configuration.
-        This method returns a PBA object based on worm's configuration.
+        Return None or False if you don't want the pba to be executed.
-        :return: An array of PBA objects.
+        :return: A pba object.
         """
-        raise NotImplementedError()
+        return self
 
-    @staticmethod
+    def should_run(self, class_name):
-    def default_get_pba(name, pba_class, linux_cmd="", windows_cmd=""):
         """
-        Default get_pba() method implementation
+        Decides if post breach action is enabled in config
-        :param name: PBA name
+        :return: True if it needs to be ran, false otherwise
-        :param pba_class: class instance. Class's name is matched to config to determine
-        if corresponding field was enabled in post breach array or not.
-        :param linux_cmd: commands for linux
-        :param windows_cmd: commands for windows
-        :return: post breach action
         """
-        if pba_class.__name__ in WormConfiguration.post_breach_actions:
+        return class_name in WormConfiguration.post_breach_actions
-            command = PBA.choose_command(linux_cmd, windows_cmd)
-            if command:
-                return PBA(name, command)
 
     def run(self):
         """
@@ -55,8 +45,7 @@ class PBA(object):
         result = exec_funct()
         ControlClient.send_telemetry('post_breach', {'command': self.command,
                                                      'result': result,
-                                                     'name': self.name,
+                                                     'name': self.name})
-                                                     'guid': GUID})
 
     def _execute_default(self):
         """

monkey/infection_monkey/post_breach/post_breach_handler.py

@@ -1,6 +1,7 @@
 import logging
 import inspect
 import importlib
+from infection_monkey.post_breach.pba import PBA
 from infection_monkey.post_breach.actions import get_pba_files
 from infection_monkey.utils import is_windows_os
 
@@ -41,10 +42,10 @@ class PostBreach(object):
             module = importlib.import_module(PATH_TO_ACTIONS + pba_file)
             # Get all classes in a module
             pba_classes = [m[1] for m in inspect.getmembers(module, inspect.isclass)
-                           if ((m[1].__module__ == module.__name__) and getattr(m[1], "get_pba", False))]
+                           if ((m[1].__module__ == module.__name__) and issubclass(m[1], PBA))]
             # Get post breach action object from class
             for pba_class in pba_classes:
-                pba = pba_class.get_pba()
+                pba = pba_class()
-                if pba:
+                if pba.should_run(pba_class.__name__):
                     pba_list.append(pba)
         return pba_list