From 21bc978f4599b40fb3dd6eaec1c79f1daf31bff1 Mon Sep 17 00:00:00 2001 From: Xiaoli Chan <2209553467@qq.com> Date: Wed, 30 Aug 2023 02:13:59 +0800 Subject: [PATCH] [connection] neff: review II Signed-off-by: Xiaoli Chan <2209553467@qq.com> --- cme/modules/rdp.py | 11 +++++------ cme/protocols/smb.py | 3 ++- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/cme/modules/rdp.py b/cme/modules/rdp.py index 28f49d99..695d2916 100644 --- a/cme/modules/rdp.py +++ b/cme/modules/rdp.py @@ -83,12 +83,10 @@ class CMEModule: context.log.fail(f"Enable RDP via smb error: {str(e)}") elif self.method == "wmi": context.log.info("Executing over WMI(ncacn_ip_tcp)") - try: - wmi_rdp = rdp_WMI(context, connection, self.dcom_timeout) - except: - pass - if "wmi_rdp" in locals() and hasattr(wmi_rdp, '_rdp_WMI__iWbemLevel1Login'): + wmi_rdp = rdp_WMI(context, connection, self.dcom_timeout) + + if hasattr(wmi_rdp, '_rdp_WMI__iWbemLevel1Login'): if "ram" in self.action: # Nt version under 6 not support RAM. try: @@ -259,7 +257,8 @@ class rdp_WMI: self.__iWbemLevel1Login = wmi.IWbemLevel1Login(iInterface) except Exception as e: self.logger.fail(f'Unexpected wmi error: {str(e)}, please try to use "-o" with "METHOD=smb"') - self.__dcom.disconnect() + if self.__iWbemLevel1Login in locals(): + self.__dcom.disconnect() def rdp_Wrapper(self, action, old=False): if old == False: diff --git a/cme/protocols/smb.py b/cme/protocols/smb.py index 1189fcbe..a9cc2265 100755 --- a/cme/protocols/smb.py +++ b/cme/protocols/smb.py @@ -1248,7 +1248,8 @@ class smb(connection): iEnumWbemClassObject = iWbemServices.ExecQuery(wmi_query) except Exception as e: self.logger.fail('Execute WQL error: {}'.format(e)) - dcom.disconnect() + if "iWbemLevel1Login" in locals(): + dcom.disconnect() else: self.logger.info(f"Executing WQL syntax: {wmi_query}") while True: