diff --git a/windows/com.py b/windows/com.py index 3581a6d..b272a40 100644 --- a/windows/com.py +++ b/windows/com.py @@ -106,6 +106,12 @@ class ImprovedVariant(VARIANT): raise ValueError("asbyte on non-VT_UI1 variant") return self._VARIANT_NAME_3.bVal + @property + def asunknown(self): + if not self.vt in [VT_UNKNOWN]: + raise ValueError("asunknown on non-VT_UNKNOWN variant") + return self._VARIANT_NAME_3.punkVal + @property def asarray(self): if not self.vt & VT_ARRAY: @@ -115,6 +121,59 @@ class ImprovedVariant(VARIANT): return ImprovedSAFEARRAY.from_PSAFEARRAY(self._VARIANT_NAME_3.parray) + @property + def aslong_array(self): + if not self.vt & VT_I4: + raise ValueError("as_bstr_array on non-VT_BSTR variant") + return self.asarray.to_list(LONG) + + def generate_asarray_property(vttype): + @property + def as_array_generated(self): + # TODO: vt check like the others ? + return self.asarray.to_list(vttype) + return as_array_generated + + asbstr_array = generate_asarray_property(BSTR) + aslong_array = generate_asarray_property(LONG) + asbyte_array = generate_asarray_property(BYTE) + asbool_array = generate_asarray_property(VARIANT_BOOL) + + def to_pyobject(self): + # if self.vt & VT_ARRAY: + # # Something better TODO i guess + # if self.vt & VT_TYPEMASK == VT_BSTR: + # import pdb;pdb.set_trace() + # print("VT_TYPEMASK ARRAY") + # return self.asarray.to_list(BSTR) + # if self.vt & VT_TYPEMASK == VT_I4: + # import pdb;pdb.set_trace() + # print("VT_TYPEMASK ARRAY") + # return self.asarray.to_list(LONG) + # raise NotImplementedError("Variant of type {0:#x}".format(self.vt)) + # use the ImprovedVariant.MAPPER that dispatch by self.vt + try: + return self.MAPPER[self.vt](self) + except KeyError: + raise NotImplementedError("Variant of type {0:#x}".format(self.vt)) + + +ImprovedVariant.MAPPER = { + VT_UI1: ImprovedVariant.asbyte.fget, + VT_I2: ImprovedVariant.asshort.fget, + VT_DISPATCH: ImprovedVariant.asdispatch.fget, + VT_BOOL: ImprovedVariant.asbool.fget, + VT_I4: ImprovedVariant.aslong.fget, + VT_BSTR: ImprovedVariant.asbstr.fget, + VT_EMPTY: (lambda x: None), + VT_NULL: (lambda x: None), + VT_UNKNOWN: ImprovedVariant.asunknown.fget, + (VT_ARRAY | VT_BSTR): ImprovedVariant.asbstr_array.fget, + (VT_ARRAY | VT_I4): ImprovedVariant.aslong_array.fget, + (VT_ARRAY | VT_UI1): ImprovedVariant.asbyte_array.fget, + (VT_ARRAY | VT_BOOL): ImprovedVariant.asbool_array.fget +} + def create_instance(clsiid, targetinterface, custom_iid=None, context=CLSCTX_INPROC_SERVER | CLSCTX_LOCAL_SERVER): if custom_iid is None: diff --git a/windows/pe_parse.py b/windows/pe_parse.py index 0c51602..1fa98f3 100644 --- a/windows/pe_parse.py +++ b/windows/pe_parse.py @@ -39,7 +39,7 @@ def get_structure_transformer_for_target(target, targetbitness=None): return ctypes_structure_transformer, create_structure_at def get_pe_bitness(baseaddr, target): - # We can force bitness as the filed we access are bitness-independant + # We can force bitness as the field we access are bitness-independant pe = GetPEFile(baseaddr, target, force_bitness=32) machine = pe.get_NT_HEADER().FileHeader.Machine if machine == 0x14c: diff --git a/windows/winobject/wmi.py b/windows/winobject/wmi.py index 9e4cf50..4db127a 100644 --- a/windows/winobject/wmi.py +++ b/windows/winobject/wmi.py @@ -8,6 +8,7 @@ from ctypes.wintypes import * import windows.com from windows.generated_def.winstructs import * from windows.generated_def.interfaces import IWbemLocator, IWbemServices, IEnumWbemClassObject, IWbemClassObject +# import windows.generated_def as gdef class WmiRequester(object): @@ -26,12 +27,12 @@ class WmiRequester(object): locator.ConnectServer(target, user, password , None, 0x80, None, None, ctypes.byref(service)) self.service = service - def select(self, frm, attrs="*"): + def select(self, frm, attrs="*", **kwargs): """Select ``attrs`` from ``frm`` :rtype: list of dict """ - return self.query("select * from {0}".format(frm), attrs) + return self.query("select * from {0}".format(frm), attrs, **kwargs) @property def classes(self): @@ -41,62 +42,78 @@ class WmiRequester(object): """ return [x["__CLASS"] for x in self.query('SELECT * FROM meta_class', attrs=["__CLASS"])] - def query(self, query, attrs="*"): + def query(self, query, attrs="*", timeout=WBEM_INFINITE): """Execute WMI ``query`` and return the attributes ``attrs`` + Timeout is not applied for the full query time but the time to retrieve one object each time. :rtype: list of dict """ + return list(self.gen_query(query, attrs, timeout)) + + def gen_query(self, query, attrs="*", timeout=WBEM_INFINITE): + """Execute WMI ``query`` and return a generator that will yield the ``attrs`` for one object each time. + Each iteration is susceptible to raise. + + :rtype: generator + """ + enumerator = self._exec_query(query, WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY) + try: + for obj, retval in self._enumerator_values_generator(enumerator, timeout=timeout): + if obj is None: + raise WindowsError(retval, WBEMSTATUS(retval & 0xffffffff).value) + yield self._iwbemclassobject_to_dict(obj, attrs) + finally: + enumerator.Release() + + def _exec_query(self, query, flags, ctx=None): enumerator = IEnumWbemClassObject() try: - self.service.ExecQuery("WQL", query, 0x20, 0, ctypes.byref(enumerator)) + self.service.ExecQuery("WQL", query, flags, ctx, ctypes.byref(enumerator)) except WindowsError as e: if (e.winerror & 0xffffffff) == WBEM_E_INVALID_CLASS: raise WindowsError(e.winerror, 'WBEM_E_INVALID_CLASS '.format(query)) elif (e.winerror & 0xffffffff) in WBEMSTATUS.values: raise WindowsError(e.winerror, WBEMSTATUS(e.winerror & 0xffffffff).value) raise + return enumerator - count = ctypes.c_ulong(0) + def _enumerator_values_generator(self, enumerator, timeout=WBEM_INFINITE): + count = ULONG(0) processor = IWbemClassObject() - res = [] - enumerator.Next(0xffffffff, 1, ctypes.byref(processor), ctypes.byref(count)) - while count.value: - current_res = {} - variant_res = windows.com.ImprovedVariant() - if attrs == "*": - attrs = [x for x in self.get_names(processor) if not x.startswith("__")] - for name in attrs: - try: - processor.Get(name, 0, ctypes.byref(variant_res), None, None) - except WindowsError as e: - if (e.winerror & 0xffffffff) == WBEM_E_NOT_FOUND: - raise WindowsError(e.winerror, 'WBEM_E_NOT_FOUND '.format(name)) - if (e.winerror & 0xffffffff) in WBEMSTATUS.values: - raise WindowsError(e.winerror, WBEMSTATUS(e.winerror & 0xffffffff).value) - raise - # TODO: something clean and generic - if variant_res.vt & VT_ARRAY: - if variant_res.vt & VT_TYPEMASK == VT_BSTR: - current_res[name] = variant_res.asarray.to_list(BSTR) - if variant_res.vt & VT_TYPEMASK == VT_I4: - current_res[name] = variant_res.asarray.to_list(LONG) - elif variant_res.vt in [VT_EMPTY, VT_NULL]: - current_res[name] = None - elif variant_res.vt == VT_BSTR: - current_res[name] = variant_res.asbstr - elif variant_res.vt == VT_I4: - current_res[name] = variant_res.aslong - elif variant_res.vt == VT_BOOL: - current_res[name] = variant_res.asbool - elif variant_res.vt == VT_I2: - current_res[name] = variant_res.asshort - elif variant_res.vt == VT_UI1: - current_res[name] = variant_res.asbyte - else: - print("[WARN] WMI Ignore variant of type {0}".format(hex(variant_res.vt))) - res.append(current_res) - enumerator.Next(0xffffffff, 1, ctypes.byref(processor), ctypes.byref(count)) - return res + result = 0 + while result != WBEM_S_FALSE: + try: + result = enumerator.Next(timeout, 1, ctypes.byref(processor), ctypes.byref(count)) + except WindowsError as e: + if (e.winerror & 0xffffffff) == WBEM_E_INVALID_CLASS: + raise WindowsError(e.winerror, 'WBEM_E_INVALID_CLASS '.format(query)) + if (e.winerror & 0xffffffff) in WBEMSTATUS.values: + raise WindowsError(e.winerror, WBEMSTATUS(e.winerror & 0xffffffff).value) + raise + procres = processor if count else None + if result != WBEM_S_FALSE: + yield procres, result + + + def _iwbemclassobject_to_dict(self, wbemclassobj, attrs): + if attrs == "*": + attrs = [x for x in self.get_names(wbemclassobj) if not x.startswith("__")] + obj_as_dict = {} + variant_res = windows.com.ImprovedVariant() + for name in attrs: + try: + wbemclassobj.Get(name, 0, ctypes.byref(variant_res), None, None) + except WindowsError as e: + if (e.winerror & 0xffffffff) == WBEM_E_NOT_FOUND: + raise WindowsError(e.winerror, 'WBEM_E_NOT_FOUND '.format(name)) + if (e.winerror & 0xffffffff) in WBEMSTATUS.values: + raise WindowsError(e.winerror, WBEMSTATUS(e.winerror & 0xffffffff).value) + raise + try: + obj_as_dict[name] = variant_res.to_pyobject() + except NotImplementedError as e: + print("[WMI-ERROR] Field <{0}> ignored: {1}".format(name, e)) + return obj_as_dict def get_names(self, processor):