Refactor privilege check logic

nxc/protocols/ssh.py

@@ -20,6 +20,8 @@ def __init__(self, args, db, host):
         self.protocol = "SSH"
         self.remote_version = "Unknown SSH Version"
         self.server_os_platform = "Linux"
+        self.shell_access = False
+        self.admin_privs = False
         self.uac = ""
         super().__init__(args, db, host)
 
@@ -116,12 +118,11 @@ def plaintext_login(self, username, password, private_key=""):
                 )
                 cred_id = self.db.add_credential("plaintext", username, password)
 
-            # Some IOT devices will not raise exception in self.conn._transport.auth_password / self.conn._transport.auth_publickey
-            # Also an early check if we are on Linux or not, as on windows only stderr and not stdout is returned ("id" is not implemented)
-            _, stdout, _ = self.conn.exec_command("id")
-            stdout = stdout.read().decode(self.args.codec, errors="ignore")
+            self.check_shell(cred_id)
 
-            self.check_privs(cred_id, stdout)
+            out = process_secret(self.password) if not self.args.key_file else f"{process_secret(self.password)} (keyfile: {self.args.key_file})"
+            display_shell_access = f"{self.uac}{self.server_os_platform}{' - Shell access!' if self.shell_access else ''}"
+            self.logger.success(f"{self.username}:{process_secret(out)} {self.mark_pwned()} {highlight(display_shell_access)}")
             return True
         except AuthenticationException as e:
             if "Private key file is encrypted" in str(e):
@@ -140,42 +141,47 @@ def plaintext_login(self, username, password, private_key=""):
             self.conn.close()
         return False
 
-    def check_privs(self, cred_id, stdout):
-        shell_access = False
+    def check_shell(self, cred_id):
         host_id = self.db.get_hosts(self.host)[0].id
 
-        # If we have stdout we know it must be linux, "id" is not implemented on Windows
-        if not stdout:
-            self.server_os_platform = "Windows"
-            _, stdout, _ = self.conn.exec_command("whoami /priv")
-            stdout = stdout.read().decode(self.args.codec, errors="ignore")
-            if "SeDebugPrivilege" in stdout:
-                self.admin_privs = True
-            elif "SeUndockPrivilege" in stdout:
-                self.admin_privs = True
-                self.uac = "with UAC - "
-
-        if not stdout:
-            self.logger.debug(f"User: {self.username} can't get a basic shell")
-            self.server_os_platform = "Network Devices"
-            shell_access = False
-        else:
-            shell_access = True
-
-        self.db.add_loggedin_relation(cred_id, host_id, shell=shell_access)
-
-        if shell_access and self.server_os_platform == "Linux":
+        # Some IOT devices will not raise exception in self.conn._transport.auth_password / self.conn._transport.auth_publickey
+        # Check Linux
+        stdout = self.conn.exec_command("id")[1].read().decode(self.args.codec, errors="ignore")
+        if stdout:
+            self.server_os_platform = "Linux"
+            self.logger.debug(f"Linux detected for user: {stdout}")
+            self.shell_access = True
             self.check_linux_priv()
             if self.admin_privs:
                 self.logger.debug(f"User {self.username} logged in successfully and is root!")
                 if self.args.key_file:
                     self.db.add_admin_user("key", self.username, self.password, host_id=host_id, cred_id=cred_id)
                 else:
                     self.db.add_admin_user("plaintext", self.username, self.password, host_id=host_id, cred_id=cred_id)
+            return
+
+        # Check Windows
+        stdout = self.conn.exec_command("whoami /priv")[1].read().decode(self.args.codec, errors="ignore")
+        if stdout:
+            self.server_os_platform = "Windows"
+            self.logger.debug(f"Windows detected for user: {stdout}")
+            self.shell_access = True
+            self.check_windows_priv(stdout)
+            self.db.add_loggedin_relation(cred_id, host_id, shell=self.shell_access)
+            return
+
+        # No shell access
+        self.shell_access = False
+        self.logger.debug(f"User: {self.username} can't get a basic shell")
+        self.server_os_platform = "Network Devices"
+        self.db.add_loggedin_relation(cred_id, host_id, shell=self.shell_access)
 
-        out = process_secret(self.password) if not self.args.key_file else f"{process_secret(self.password)} (keyfile: {self.args.key_file})"
-        display_shell_access = f"{self.uac}{self.server_os_platform}{' - Shell access!' if shell_access else ''}"
-        self.logger.success(f"{self.username}:{process_secret(out)} {self.mark_pwned()} {highlight(display_shell_access)}")
+    def check_windows_priv(self, stdout):
+        if "SeDebugPrivilege" in stdout:
+            self.admin_privs = True
+        elif "SeUndockPrivilege" in stdout:
+            self.admin_privs = True
+            self.uac = "with UAC - "
 
     def check_linux_priv(self):
         self.admin_privs = False