Improving com / wmi API / code

TODO

@@ -36,6 +36,8 @@ TODO:
 
         * What about ``NtWow64QueryVirtualMemory64`` ?
 
+    - Parse .IDL file for more COM NAME->IID
+
 
 CHANGELOG:
 

windows/__init__.py

@@ -11,6 +11,7 @@ Exported:
     current_thread : :class:`windows.winobject.CurrentThread`
 """
 
+
 from windows import winproxy
 from windows import winobject
 
@@ -33,6 +34,7 @@ import windows.utils
 import windows.debug
 import windows.wintrust
 import windows.syswow64
+import windows.com
 
 __all__ = ["system", 'current_process', 'current_thread']
 

windows/com.py

@@ -1,7 +1,7 @@
 import struct
 import ctypes
 import functools
-from ctypes.wintypes import HRESULT, byref, pointer
+from ctypes.wintypes import HRESULT, byref, pointer, cast
 
 import windows
 from windows import winproxy
@@ -30,6 +30,51 @@ def init():
     return winproxy.CoInitializeSecurity(0, -1, None, 0, 0, RPC_C_IMP_LEVEL_IMPERSONATE, 0,0,0)
 
 
+class ImprovedSAFEARRAY(SAFEARRAY):
+        @classmethod
+        def of_type(cls, addr, t):
+            self = cls.from_address(addr)
+            self.elt_type = t
+            return self
+
+        @classmethod
+        def from_PSAFEARRAY(self, psafearray):
+            res = cast(psafearray, POINTER(ImprovedSAFEARRAY))[0]
+            return res
+
+        def to_list(self, t=None):
+            if t is None:
+                if hasattr(self, "elt_type"):
+                    t = self.elt_type
+                else:
+                    raise ValueError("Missing type of the array")
+            if self.cDims !=  1:
+                raise NotImplementedError("tagSAFEARRAY if dims != 1")
+
+            nb_element = self.rgsabound[0].cElements
+            llbound = self.rgsabound[0].lLbound
+            if self.cbElements != ctypes.sizeof(t):
+                raise ValueError("Size of elements != sizeof(type)")
+            data = [t.from_address(self.pvData + (i + llbound) * ctypes.sizeof(t)).value for i in range(nb_element)]
+            return data
+
+#VT_VALUE_TO_TYPE = {
+#VT_I2 : SHORT,
+#VT_I4 : LONG,
+#VT_BSTR : BSTR,
+#VT_VARIANT : VARIANT,
+#VT_UI1 : UCHAR,
+#VT_UI2 : USHORT,
+#VT_UI4 : DWORD,
+#VT_I8 : LONGLONG,
+#VT_UI8 : ULONG64,
+#VT_INT : INT,
+#VT_UINT : UINT,
+#VT_HRESULT : HRESULT,
+#VT_PTR : PVOID,
+#VT_LPSTR : LPCSTR,
+#VT_LPWSTR : LPWSTR,
+#}
 
 class ImprovedVariant(VARIANT):
     @property
@@ -57,6 +102,26 @@ class ImprovedVariant(VARIANT):
             raise ValueError("asdispatch on non-VT_DISPATCH variant")
         return interfaces.IDispatch(self._VARIANT_NAME_3.pdispVal)
 
+    @property
+    def asshort(self):
+        if not self.vt in [VT_I2]:
+            raise ValueError("asshort on non-VT_I2 variant")
+        return self._VARIANT_NAME_3.iVal
+
+    @property
+    def asbyte(self):
+        if not self.vt in [VT_UI1]:
+            raise ValueError("asbyte on non-VT_UI1 variant")
+        return self._VARIANT_NAME_3.bVal
+
+    @property
+    def asarray(self):
+        if not self.vt & VT_ARRAY:
+            raise ValueError("asarray on non-VT_ARRAY variant")
+        # TODO: auto extract VT_TYPE for the array ?
+        #type = VT_VALUE_TO_TYPE[self.vt & VT_TYPEMASK]
+        return ImprovedSAFEARRAY.from_PSAFEARRAY(self._VARIANT_NAME_3.parray)
+
 
 
 def create_instance(clsiid, targetinterface, custom_iid=None):

windows/winobject/wmi.py

@@ -10,28 +10,7 @@ from windows.generated_def.winstructs import *
 from windows.generated_def.interfaces import IWbemLocator, IWbemServices, IEnumWbemClassObject, IWbemClassObject
 
 
-class ImprovedSAFEARRAY(windows.generated_def.winstructs.SAFEARRAY):
-        @classmethod
-        def of_type(cls, addr, t):
-            self = cls.from_address(addr)
-            self.elt_type = t
-            return self
 
-        def to_list(self, t=None):
-            if t is None:
-                if hasattr(self, "elt_type"):
-                    t = self.elt_type
-                else:
-                    raise ValueError("Missing type of the array")
-            if self.cDims !=  1:
-                raise NotImplementedError("tagSAFEARRAY if dims != 1")
-
-            nb_element = self.rgsabound[0].cElements
-            llbound = self.rgsabound[0].lLbound
-            if self.cbElements != ctypes.sizeof(t):
-                raise ValueError("Size of elements != sizeof(type)")
-            data = [t.from_address(self.pvData + (i + llbound) * ctypes.sizeof(t)).value for i in range(nb_element)]
-            return data
 
 
 
@@ -84,9 +63,9 @@ class WmiRequester(object):
                 # TODO: something clean and generic
                 if variant_res.vt & VT_ARRAY:
                     if variant_res.vt & VT_TYPEMASK == VT_BSTR:
-                        current_res[name] = variant_res._Data.parray[0].to_list(BSTR)
+                        current_res[name] = variant_res.asarray.to_list(BSTR)
                     if variant_res.vt & VT_TYPEMASK == VT_I4:
-                        current_res[name] = variant_res._Data.parray[0].to_list(LONG)
+                        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:
@@ -95,8 +74,12 @@ class WmiRequester(object):
                     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("Ignore variant of type {0}".format(hex(variant_res.vt)))
+                    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
@@ -104,6 +87,6 @@ class WmiRequester(object):
     def get_names(self, processor):
         res = POINTER(SAFEARRAY)()
         processor.GetNames(None, 0, None, byref(res))
-        safe_array = ctypes.cast(res, POINTER(ImprovedSAFEARRAY))[0]
+        safe_array = ctypes.cast(res, POINTER(windows.com.ImprovedSAFEARRAY))[0]
         safe_array.elt_type = BSTR
         return safe_array.to_list()