Refactored new user in linux to AutoNewLinuxUser and created AutoNewUser ABC

monkey/infection_monkey/post_breach/actions/communicate_as_new_user.py

@@ -5,7 +5,7 @@ import string
 import subprocess
 import time
 
-from infection_monkey.utils.windows.auto_new_user import AutoNewUser, NewUserError
+from infection_monkey.utils.windows.auto_new_user import NewUserError, create_auto_new_user
 from common.data.post_breach_consts import POST_BREACH_COMMUNICATE_AS_NEW_USER
 from infection_monkey.post_breach.pba import PBA
 from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
@@ -47,19 +47,12 @@ class CommunicateAsNewUser(PBA):
 
     def communicate_as_new_user_linux(self, username):
         try:
-            # add user + ping
-            linux_cmds = get_linux_commands_to_add_user(username)
-            commandline = "ping -c 1 {}".format(PING_TEST_DOMAIN)
-            linux_cmds.extend([";", "sudo", "-u", username, commandline])
-            final_command = ' '.join(linux_cmds)
-            exit_status = os.system(final_command)
-            self.send_ping_result_telemetry(exit_status, commandline, username)
-            # delete the user.
-            commands_to_delete_user = get_linux_commands_to_delete_user(username)
-            logger.debug("Trying to delete the user {} with commands {}".format(username, str(commands_to_delete_user)))
-            delete_user_output = subprocess.check_output(" ".join(commands_to_delete_user), stderr=subprocess.STDOUT, shell=True)
-            logger.debug("Deletion output: {}".format(delete_user_output))
-            # Leaking the process on purpose - nothing we can do if it's stuck.
+            with create_auto_new_user(username, PASSWORD) as new_user:
+                commandline = "sudo -u {username} ping -c 1 {domain}".format(
+                    username=new_user.username,
+                    domain=PING_TEST_DOMAIN)
+                exit_status = os.system(commandline)
+                self.send_ping_result_telemetry(exit_status, commandline, new_user.username)
         except subprocess.CalledProcessError as e:
             PostBreachTelem(self, (e.output, False)).send()
 
@@ -71,7 +64,7 @@ class CommunicateAsNewUser(PBA):
         import win32event
 
         try:
-            with AutoNewUser(username, PASSWORD) as new_user:
+            with create_auto_new_user(username, PASSWORD) as new_user:
                 # Using os.path is OK, as this is on windows for sure
                 ping_app_path = os.path.join(os.environ["WINDIR"], "system32", "PING.exe")
                 if not os.path.exists(ping_app_path):

monkey/infection_monkey/utils/windows/auto_new_user.py

@@ -1,7 +1,9 @@
 import logging
 import subprocess
+import abc
 
-from infection_monkey.post_breach.actions.add_user import BackdoorUser
+from infection_monkey.utils.environment import is_windows_os
+from infection_monkey.utils.linux.users import get_linux_commands_to_add_user, get_linux_commands_to_delete_user
 from infection_monkey.utils.windows.users import get_windows_commands_to_delete_user, get_windows_commands_to_add_user, \
     get_windows_commands_to_deactivate_user
 
@@ -12,26 +14,79 @@ class NewUserError(Exception):
     pass
 
 
-class AutoNewUser(object):
+class AutoNewUser:
     """
-    RAII object to use for creating and using a new user in Windows. Use with `with`.
+    RAII object to use for creating and using a new user. Use with `with`.
     User will be created when the instance is instantiated.
-    User will log on at the start of the `with` scope.
-    User will log off and get deleted at the end of said `with` scope.
+    User will be available for use (log on for Windows, for example) at the start of the `with` scope.
+    User will be removed (deactivated and deleted for Windows, for example) at the end of said `with` scope.
 
     Example:
-             # Created                           # Logged on
-        with AutoNewUser("user", "pass") as new_user:
+             # Created                                                 # Logged on
+        with AutoNewUser("user", "pass", is_on_windows()) as new_user:
             ...
             ...
         # Logged off and deleted
         ...
+        """
+    __metaclass__ = abc.ABCMeta
+
+    def __init__(self, username, password):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def __enter__(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        raise NotImplementedError()
+
+
+def create_auto_new_user(username, password, is_windows=is_windows_os()):
+    if is_windows:
+        return AutoNewWindowsUser(username, password)
+    else:
+        return AutoNewLinuxUser(username, password)
+
+
+class AutoNewLinuxUser(AutoNewUser):
+    """
+    See AutoNewUser's documentation for details.
     """
     def __init__(self, username, password):
         """
         Creates a user with the username + password.
         :raises: subprocess.CalledProcessError if failed to add the user.
         """
+        super(AutoNewLinuxUser, self).__init__(username, password)
+        self.username = username
+        self.password = password
+
+        commands_to_add_user = get_linux_commands_to_add_user(username)
+        logger.debug("Trying to add {} with commands {}".format(self.username, str(commands_to_add_user)))
+        _ = subprocess.check_output(' '.join(commands_to_add_user), stderr=subprocess.STDOUT, shell=True)
+
+    def __enter__(self):
+        pass  # No initialization/logging on needed in Linux
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # delete the user.
+        commands_to_delete_user = get_linux_commands_to_delete_user(self.username)
+        logger.debug("Trying to delete {} with commands {}".format(self.username, str(commands_to_delete_user)))
+        _ = subprocess.check_output(" ".join(commands_to_delete_user), stderr=subprocess.STDOUT, shell=True)
+
+
+class AutoNewWindowsUser(AutoNewUser):
+    """
+    See AutoNewUser's documentation for details.
+    """
+    def __init__(self, username, password):
+        """
+        Creates a user with the username + password.
+        :raises: subprocess.CalledProcessError if failed to add the user.
+        """
+        super(AutoNewWindowsUser, self).__init__(username, password)
         self.username = username
         self.password = password
 
@@ -65,14 +120,7 @@ class AutoNewUser(object):
 
         # Try to disable and then delete the user.
         self.try_deactivate_user()
-        self.try_disable_user()
-
-    def try_disable_user(self):
-        try:
-            _ = subprocess.check_output(
-                get_windows_commands_to_delete_user(self.username), stderr=subprocess.STDOUT, shell=True)
-        except Exception as err:
-            raise NewUserError("Can't delete user {}. Info: {}".format(self.username, err))
+        self.try_delete_user()
 
     def try_deactivate_user(self):
         try:
@@ -80,3 +128,10 @@ class AutoNewUser(object):
                 get_windows_commands_to_deactivate_user(self.username), stderr=subprocess.STDOUT, shell=True)
         except Exception as err:
             raise NewUserError("Can't deactivate user {}. Info: {}".format(self.username, err))
+
+    def try_delete_user(self):
+        try:
+            _ = subprocess.check_output(
+                get_windows_commands_to_delete_user(self.username), stderr=subprocess.STDOUT, shell=True)
+        except Exception as err:
+            raise NewUserError("Can't delete user {}. Info: {}".format(self.username, err))