Moved runas to autonewuser class and deleted duplication from communicate_as_new_user

monkey/infection_monkey/post_breach/actions/communicate_as_new_user.py

@@ -13,12 +13,10 @@ from infection_monkey.utils.environment import is_windows_os
 
 PING_TEST_DOMAIN = "google.com"
 
-PING_WAIT_TIMEOUT_IN_MILLISECONDS = 20 * 1000
-
 CREATED_PROCESS_AS_USER_PING_SUCCESS_FORMAT = "Created process '{}' as user '{}', and successfully pinged."
 CREATED_PROCESS_AS_USER_PING_FAILED_FORMAT = "Created process '{}' as user '{}', but failed to ping (exit status {})."
 
-USERNAME = "somenewuser"
+USERNAME_PREFIX = "somenewuser"
 PASSWORD = "N3WPa55W0rD!1"
 
 logger = logging.getLogger(__name__)
@@ -35,90 +33,26 @@ class CommunicateAsNewUser(PBA):
 
     def run(self):
         username = CommunicateAsNewUser.get_random_new_user_name()
-        if is_windows_os():
+        try:
-            self.communicate_as_new_user_windows(username)
+            with create_auto_new_user(username, PASSWORD) as new_user:
-        else:
+                ping_commandline = CommunicateAsNewUser.get_commandline_for_ping()
-            self.communicate_as_new_user_linux(username)
+                exit_status = new_user.run_as(ping_commandline)
+                self.send_ping_result_telemetry(exit_status, ping_commandline, username)
+        except subprocess.CalledProcessError as e:
+            PostBreachTelem(self, (e.output, False)).send()
+        except NewUserError as e:
+            PostBreachTelem(self, (str(e), False)).send()
 
     @staticmethod
     def get_random_new_user_name():
-        return USERNAME + ''.join(random.choice(string.ascii_lowercase) for _ in range(5))
+        return USERNAME_PREFIX + ''.join(random.choice(string.ascii_lowercase) for _ in range(5))
 
-    def communicate_as_new_user_linux(self, username):
+    @staticmethod
-        try:
+    def get_commandline_for_ping(domain=PING_TEST_DOMAIN, is_windows=is_windows_os()):
-            with create_auto_new_user(username, PASSWORD, is_windows=False) as _:
+        if is_windows:
-                commandline = "sudo -u {username} ping -c 1 {domain}".format(
+            return "{} {} {} {}".format("PING.exe", domain, "-n", "1")
-                    username=username,
+        else:
-                    domain=PING_TEST_DOMAIN)
+            return "ping -c 1 {domain}".format(domain=PING_TEST_DOMAIN)
-                exit_status = os.system(commandline)
-                self.send_ping_result_telemetry(exit_status, commandline, username)
-        except subprocess.CalledProcessError as e:
-            PostBreachTelem(self, (e.output, False)).send()
-
-    def communicate_as_new_user_windows(self, username):
-        # Importing these only on windows, as they won't exist on linux.
-        import win32con
-        import win32process
-        import win32api
-        import win32event
-
-        try:
-            with create_auto_new_user(username, PASSWORD, is_windows=True) 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):
-                    PostBreachTelem(self, ("{} not found.".format(ping_app_path), False)).send()
-                    return  # Can't continue without ping.
-
-                try:
-                    # Open process as that user:
-                    # https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-createprocessasusera
-                    commandline = "{} {} {} {}".format(ping_app_path, PING_TEST_DOMAIN, "-n", "1")
-                    process_handle, thread_handle, _, _ = win32process.CreateProcessAsUser(
-                        new_user.get_logon_handle(),  # A handle to the primary token that represents a user.
-                        None,  # The name of the module to be executed.
-                        commandline,  # The command line to be executed.
-                        None,  # Process attributes
-                        None,  # Thread attributes
-                        True,  # Should inherit handles
-                        win32con.NORMAL_PRIORITY_CLASS,  # The priority class and the creation of the process.
-                        None,  # An environment block for the new process. If this parameter is NULL, the new process
-                        # uses the environment of the calling process.
-                        None,  # CWD. If this parameter is NULL, the new process will have the same current drive and
-                        # directory as the calling process.
-                        win32process.STARTUPINFO()  # STARTUPINFO structure.
-                        # https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/ns-processthreadsapi-startupinfoa
-                    )
-
-                    logger.debug(
-                        "Waiting for ping process to finish. Timeout: {}ms".format(PING_WAIT_TIMEOUT_IN_MILLISECONDS))
-
-                    # Ignoring return code, as we'll use `GetExitCode` to determine the state of the process later.
-                    _ = win32event.WaitForSingleObject(  # Waits until the specified object is signaled, or time-out.
-                        process_handle,  # Ping process handle
-                        PING_WAIT_TIMEOUT_IN_MILLISECONDS  # Timeout in milliseconds
-                    )
-
-                    ping_exit_code = win32process.GetExitCodeProcess(process_handle)
-
-                    self.send_ping_result_telemetry(ping_exit_code, commandline, username)
-                except Exception as e:
-                    # If failed on 1314, it's possible to try to elevate the rights of the current user with the
-                    #  "Replace a process level token" right, using Local Security Policy editing.
-                    PostBreachTelem(self, (
-                        "Failed to open process as user {}. Error: {}".format(username, str(e)), False)).send()
-                finally:
-                    try:
-                        win32api.CloseHandle(process_handle)
-                        win32api.CloseHandle(thread_handle)
-                    except Exception as err:
-                        logger.error("Close handle error: " + str(err))
-        except subprocess.CalledProcessError as err:
-            PostBreachTelem(self, (
-                "Couldn't create the user '{}'. Error output is: '{}'".format(username, str(err)),
-                False)).send()
-        except NewUserError as e:
-            PostBreachTelem(self, (str(e), False)).send()
 
     def send_ping_result_telemetry(self, exit_status, commandline, username):
         """

monkey/infection_monkey/utils/auto_new_user.py

@@ -32,3 +32,7 @@ class AutoNewUser:
     @abc.abstractmethod
     def __exit__(self, exc_type, exc_val, exc_tb):
         raise NotImplementedError()
+
+    @abc.abstractmethod
+    def run_as(self, command):
+        raise NotImplementedError()

monkey/infection_monkey/utils/linux/users.py

@@ -1,5 +1,6 @@
 import datetime
 import logging
+import os
 import subprocess
 
 from infection_monkey.utils.auto_new_user import AutoNewUser
@@ -31,6 +32,7 @@ class AutoNewLinuxUser(AutoNewUser):
     """
     See AutoNewUser's documentation for details.
     """
+
     def __init__(self, username, password):
         """
         Creates a user with the username + password.
@@ -45,6 +47,10 @@ class AutoNewLinuxUser(AutoNewUser):
     def __enter__(self):
         pass  # No initialization/logging on needed in Linux
 
+    def run_as(self, command):
+        command_as_new_user = "sudo -u {username} {command}".format(username=self.username, command=command)
+        return os.system(command_as_new_user)
+
     def __exit__(self, exc_type, exc_val, exc_tb):
         # delete the user.
         commands_to_delete_user = get_linux_commands_to_delete_user(self.username)

monkey/infection_monkey/utils/windows/users.py

@@ -5,6 +5,7 @@ from infection_monkey.utils.auto_new_user import AutoNewUser
 from infection_monkey.utils.new_user_error import NewUserError
 
 ACTIVE_NO_NET_USER = '/ACTIVE:NO'
+WAIT_TIMEOUT_IN_MILLISECONDS = 20 * 1000
 
 logger = logging.getLogger(__name__)
 
@@ -40,6 +41,58 @@ class AutoNewWindowsUser(AutoNewUser):
     """
     See AutoNewUser's documentation for details.
     """
+
+    def run_as(self, command):
+        # Importing these only on windows, as they won't exist on linux.
+        import win32con
+        import win32process
+        import win32api
+        import win32event
+
+        exit_code = -1
+        process_handle = None
+        thread_handle = None
+
+        try:
+            # Open process as that user:
+            # https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-createprocessasusera
+            process_handle, thread_handle, _, _ = win32process.CreateProcessAsUser(
+                self.get_logon_handle(),  # A handle to the primary token that represents a user.
+                None,  # The name of the module to be executed.
+                command,  # The command line to be executed.
+                None,  # Process attributes
+                None,  # Thread attributes
+                True,  # Should inherit handles
+                win32con.NORMAL_PRIORITY_CLASS,  # The priority class and the creation of the process.
+                None,  # An environment block for the new process. If this parameter is NULL, the new process
+                # uses the environment of the calling process.
+                None,  # CWD. If this parameter is NULL, the new process will have the same current drive and
+                # directory as the calling process.
+                win32process.STARTUPINFO()  # STARTUPINFO structure.
+                # https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/ns-processthreadsapi-startupinfoa
+            )
+
+            logger.debug(
+                "Waiting for process to finish. Timeout: {}ms".format(WAIT_TIMEOUT_IN_MILLISECONDS))
+
+            # Ignoring return code, as we'll use `GetExitCode` to determine the state of the process later.
+            _ = win32event.WaitForSingleObject(  # Waits until the specified object is signaled, or time-out.
+                process_handle,  # Ping process handle
+                WAIT_TIMEOUT_IN_MILLISECONDS  # Timeout in milliseconds
+            )
+
+            exit_code = win32process.GetExitCodeProcess(process_handle)
+        finally:
+            try:
+                if process_handle is not None:
+                    win32api.CloseHandle(process_handle)
+                if thread_handle is not None:
+                    win32api.CloseHandle(thread_handle)
+            except Exception as err:
+                logger.error("Close handle error: " + str(err))
+
+        return exit_code
+
     def __init__(self, username, password):
         """
         Creates a user with the username + password.
@@ -94,7 +147,7 @@ class AutoNewWindowsUser(AutoNewUser):
         try:
             commands_to_delete_user = get_windows_commands_to_delete_user(self.username)
             logger.debug(
-                "Trying to deactivate {} with commands {}".format(self.username, str(commands_to_delete_user)))
+                "Trying to delete {} with commands {}".format(self.username, str(commands_to_delete_user)))
             _ = subprocess.check_output(
                 commands_to_delete_user, stderr=subprocess.STDOUT, shell=True)
         except Exception as err: