From 20459ef3fda9d1ecc0b394ea70c281012ec055af Mon Sep 17 00:00:00 2001
From: XiaoliChan <2209553467@qq.com>
Date: Thu, 14 Sep 2023 11:14:30 +0800
Subject: [PATCH] [wmi] Revert new check_admin function

Signed-off-by: XiaoliChan <2209553467@qq.com>
---
 cme/connection.py    |  6 +-----
 cme/protocols/wmi.py | 41 ++++++++++++++++++++++++-----------------
 2 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/cme/connection.py b/cme/connection.py
index 5c86b196e..09aa0ef39 100755
--- a/cme/connection.py
+++ b/cme/connection.py
@@ -42,7 +42,7 @@ def _decorator(self, *args, **kwargs):
 
     return wraps(func)(_decorator)
 
-def dcom_FirewallChecker(iInterface, timeout=None):
+def dcom_FirewallChecker(iInterface, timeout):
     stringBindings = iInterface.get_cinstance().get_string_bindings()
     for strBinding in stringBindings:
         if strBinding['wTowerId'] == 7:
@@ -60,10 +60,6 @@ def dcom_FirewallChecker(iInterface, timeout=None):
                 stringBinding = 'ncacn_ip_tcp:%s%s' % (iInterface.get_target(), bindingPort)
     if "stringBinding" not in locals():
         return True, None
-    
-    # if not timeout, which means not doing firewall check.
-    if not timeout:
-        return True, stringBinding
     try:
         rpctransport = transport.DCERPCTransportFactory(stringBinding)
         rpctransport.set_connect_timeout(timeout)
diff --git a/cme/protocols/wmi.py b/cme/protocols/wmi.py
index 72e2a5a69..dc38ecac9 100644
--- a/cme/protocols/wmi.py
+++ b/cme/protocols/wmi.py
@@ -164,29 +164,36 @@ def print_host_info(self):
     def check_if_admin(self):
         try:
             dcom = DCOMConnection(self.conn.getRemoteName(), self.username, self.password, self.domain, self.lmhash, self.nthash, oxidResolver=True, doKerberos=self.doKerberos ,kdcHost=self.kdcHost, aesKey=self.aesKey)
-            dcom.set_connect_timeout(self.args.rpc_timeout)
             iInterface = dcom.CoCreateInstanceEx(CLSID_WbemLevel1Login, IID_IWbemLevel1Login)
-            _, self.stringBinding = dcom_FirewallChecker(iInterface)
-            iWbemLevel1Login = IWbemLevel1Login(iInterface)
-            iWbemServices = iWbemLevel1Login.NTLMLogin('//./root/cimv2', NULL, NULL)
+            flag, self.stringBinding = dcom_FirewallChecker(iInterface, self.args.rpc_timeout)
         except Exception as e:
             if "dcom" in locals():
                 dcom.disconnect()
 
-            error_msg = str(e).lower()
-
-            if error_msg.find("access_denied") > 0:
-                pass
-            else:
-                if error_msg.find("timed out") > 0 or error_msg.lower().find("connection refused") > 0:
-                    error_msg = f'Check admin error: dcom initialization failed with stringbinding: "{self.stringBinding}", please try "--rpc-timeout" option. (probably is admin)'
-                elif not self.stringBinding:
+            if not str(e).lower().find("access_denied") >=0:
+                self.logger.fail(str(e))
+        else:
+            if not flag or not self.stringBinding:
+                dcom.disconnect()
+                error_msg = f'Check admin error: dcom initialization failed with stringbinding: "{self.stringBinding}", please try "--rpc-timeout" option. (probably is admin)'
+                
+                if not self.stringBinding:
                     error_msg = "Check admin error: dcom initialization failed: can't get target stringbinding, maybe cause by IPv6 or any other issues, please check your target again"
-                self.logger.fail(error_msg)
-        else:            
-            dcom.disconnect()
-            self.logger.extra['protocol'] = "WMI"
-            self.admin_privs = True
+                
+                self.logger.fail(error_msg) if not flag else self.logger.debug(error_msg)
+            else:
+                try:
+                    iWbemLevel1Login = IWbemLevel1Login(iInterface)
+                    iWbemServices = iWbemLevel1Login.NTLMLogin('//./root/cimv2', NULL, NULL)
+                except Exception as e:
+                    dcom.disconnect()
+
+                    if not str(e).lower().find("access_denied") >=0:
+                        self.logger.fail(str(e))
+                else:
+                    dcom.disconnect()
+                    self.logger.extra['protocol'] = "WMI"
+                    self.admin_privs = True
         return
 
     def kerberos_login(self, domain, username, password="", ntlm_hash="", aesKey="", kdcHost="", useCache=False):