psexec/wmi: Use context manager, reuse code

This commit is contained in:
Oleksii Shevchuk 2019-05-12 19:18:25 +03:00
parent 76da35de09
commit b7a64297ca
1 changed files with 74 additions and 82 deletions

View File

@ -15,6 +15,8 @@ except ImportError:
import encodings
from contextlib import contextmanager
from StringIO import StringIO
from base64 import b64encode
from hashlib import md5
@ -201,6 +203,40 @@ class ConnectionInfo(object):
self.aes, self.TGT, self.TGS
]
def create_dcom_connection(self):
return DCOMConnection(
self.host, self.user, self.password, self.domain,
self.lm, self.nt, self.aes, oxidResolver=True,
doKerberos=self.kerberos
)
@contextmanager
def create_wbem(self, namespace='//./root/cimv2', rpc_auth_level=None):
dcom = self.create_dcom_connection()
try:
iInterface = dcom.CoCreateInstanceEx(
wmi.CLSID_WbemLevel1Login,wmi.IID_IWbemLevel1Login)
iWbemLevel1Login = wmi.IWbemLevel1Login(iInterface)
try:
iWbemServices = iWbemLevel1Login.NTLMLogin(namespace, NULL, NULL)
if rpc_auth_level == 'privacy':
iWbemServices.get_dce_rpc().set_auth_level(
RPC_C_AUTHN_LEVEL_PKT_PRIVACY)
elif rpc_auth_level == 'integrity':
iWbemServices.get_dce_rpc().set_auth_level(
RPC_C_AUTHN_LEVEL_PKT_INTEGRITY)
yield iWbemServices
finally:
iWbemLevel1Login.RemRelease()
finally:
dcom.disconnect()
def create_connection(self, klass=SMBConnection):
try:
smb = klass(self.host, self.host, None, self.port, timeout=self.timeout)
@ -578,24 +614,9 @@ def sc(conninfo, command, output=True, wait=30):
return rpctransport.get_smb_connection(), output_filename
def wmiexec(conninfo, command, share='C$', output=True):
dcom = DCOMConnection(
conninfo.host,
conninfo.user,
conninfo.password,
conninfo.domain,
conninfo.lm, conninfo.nt,
conninfo.aes,
oxidResolver=True,
doKerberos=conninfo.kerberos
)
iInterface = dcom.CoCreateInstanceEx(wmi.CLSID_WbemLevel1Login,wmi.IID_IWbemLevel1Login)
iWbemLevel1Login = wmi.IWbemLevel1Login(iInterface)
iWbemServices= iWbemLevel1Login.NTLMLogin('//./root/cimv2', NULL, NULL)
iWbemLevel1Login.RemRelease()
output_filename = None
with conninfo.create_wbem() as iWbemServices:
win32Process, _ = iWbemServices.GetObject('Win32_Process')
if output:
output_filename = '\\'.join([
@ -607,7 +628,6 @@ def wmiexec(conninfo, command, share='C$', output=True):
)
win32Process.Create(command, r'C:\Windows\Temp', None)
dcom.disconnect()
return output_filename
@ -619,29 +639,7 @@ def wql(
host, port, user, domain, password, ntlm, timeout=timeout
)
dcom = DCOMConnection(
conninfo.host,
conninfo.user,
conninfo.password,
conninfo.domain,
conninfo.lm, conninfo.nt,
conninfo.aes,
oxidResolver=True,
doKerberos=conninfo.kerberos
)
try:
iInterface = dcom.CoCreateInstanceEx(wmi.CLSID_WbemLevel1Login,wmi.IID_IWbemLevel1Login)
iWbemLevel1Login = wmi.IWbemLevel1Login(iInterface)
try:
iWbemServices = iWbemLevel1Login.NTLMLogin(namespace, NULL, NULL)
if rpc_auth_level == 'privacy':
iWbemServices.get_dce_rpc().set_auth_level(RPC_C_AUTHN_LEVEL_PKT_PRIVACY)
elif rpc_auth_level == 'integrity':
iWbemServices.get_dce_rpc().set_auth_level(RPC_C_AUTHN_LEVEL_PKT_INTEGRITY)
with conninfo.create_wbem(namespace, rpc_auth_level) as iWbemServices:
iEnumWbemClassObject = iWbemServices.ExecQuery(query.strip())
first = True
@ -679,12 +677,6 @@ def wql(
finally:
iEnumWbemClassObject.RemRelease()
finally:
iWbemLevel1Login.RemRelease()
finally:
dcom.disconnect()
def check(host, port, user, domain, password, ntlm, timeout=30):
conninfo = ConnectionInfo(
host, port, user, domain, password, ntlm, timeout=timeout