From 90be53e920d81185a897ccb0355336214d8f3384 Mon Sep 17 00:00:00 2001
From: Shay Nehmad <shay.nehmad@guardicore.com>
Date: Thu, 3 Oct 2019 16:58:32 +0300
Subject: [PATCH] Forgot to return the instance upon __enter__ call on
 LinuxNewUser

---
 monkey/infection_monkey/utils/linux/users.py  |  2 +-
 .../infection_monkey/utils/windows/users.py   | 57 +++++++++----------
 2 files changed, 29 insertions(+), 30 deletions(-)

diff --git a/monkey/infection_monkey/utils/linux/users.py b/monkey/infection_monkey/utils/linux/users.py
index 46bc0ed87..e4968c522 100644
--- a/monkey/infection_monkey/utils/linux/users.py
+++ b/monkey/infection_monkey/utils/linux/users.py
@@ -45,7 +45,7 @@ class AutoNewLinuxUser(AutoNewUser):
         _ = subprocess.check_output(' '.join(commands_to_add_user), stderr=subprocess.STDOUT, shell=True)
 
     def __enter__(self):
-        pass  # No initialization/logging on needed in Linux
+        return self  # 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)
diff --git a/monkey/infection_monkey/utils/windows/users.py b/monkey/infection_monkey/utils/windows/users.py
index 1bc3f2738..857cddd07 100644
--- a/monkey/infection_monkey/utils/windows/users.py
+++ b/monkey/infection_monkey/utils/windows/users.py
@@ -41,6 +41,34 @@ 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)
+
+        windows_cmds = get_windows_commands_to_add_user(self.username, self.password, True)
+        logger.debug("Trying to add {} with commands {}".format(self.username, str(windows_cmds)))
+        _ = subprocess.check_output(windows_cmds, stderr=subprocess.STDOUT, shell=True)
+
+    def __enter__(self):
+        # Importing these only on windows, as they won't exist on linux.
+        import win32security
+        import win32con
+
+        try:
+            # Logon as new user: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-logonusera
+            self.logon_handle = win32security.LogonUser(
+                self.username,
+                ".",  # Use current domain.
+                self.password,
+                win32con.LOGON32_LOGON_INTERACTIVE,  # Logon type - interactive (normal user). Need this to open ping
+                                                     # using a shell.
+                win32con.LOGON32_PROVIDER_DEFAULT)  # Which logon provider to use - whatever Windows offers.
+        except Exception as err:
+            raise NewUserError("Can't logon as {}. Error: {}".format(self.username, str(err)))
+        return self
 
     def run_as(self, command):
         # Importing these only on windows, as they won't exist on linux.
@@ -93,35 +121,6 @@ class AutoNewWindowsUser(AutoNewUser):
 
         return exit_code
 
-    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)
-
-        windows_cmds = get_windows_commands_to_add_user(self.username, self.password, True)
-        logger.debug("Trying to add {} with commands {}".format(self.username, str(windows_cmds)))
-        _ = subprocess.check_output(windows_cmds, stderr=subprocess.STDOUT, shell=True)
-
-    def __enter__(self):
-        # Importing these only on windows, as they won't exist on linux.
-        import win32security
-        import win32con
-
-        try:
-            # Logon as new user: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-logonusera
-            self.logon_handle = win32security.LogonUser(
-                self.username,
-                ".",  # Use current domain.
-                self.password,
-                win32con.LOGON32_LOGON_INTERACTIVE,  # Logon type - interactive (normal user). Need this to open ping
-                                                     # using a shell.
-                win32con.LOGON32_PROVIDER_DEFAULT)  # Which logon provider to use - whatever Windows offers.
-        except Exception as err:
-            raise NewUserError("Can't logon as {}. Error: {}".format(self.username, str(err)))
-        return self
-
     def get_logon_handle(self):
         return self.logon_handle