From 59621034f550ff7e7ff934c4cc778f48629c3a9b Mon Sep 17 00:00:00 2001
From: Clement Rouault <clement.rouault@sogeti.com>
Date: Fri, 11 Dec 2015 09:54:43 +0100
Subject: [PATCH] Add some play-code for wintrust/registry/wmi

---
 registry.py |  82 +++++++++++++++++++++
 wintrust.py |  69 ++++++++++++++++++
 wmi.py      | 207 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 358 insertions(+)
 create mode 100644 registry.py
 create mode 100644 wintrust.py
 create mode 100644 wmi.py

diff --git a/registry.py b/registry.py
new file mode 100644
index 0000000..f5c3a34
--- /dev/null
+++ b/registry.py
@@ -0,0 +1,82 @@
+import _winreg
+import windows
+import itertools
+import collections
+
+
+class ExpectWindowsError(object):
+    def __init__(self, errornumber):
+        self.errornumber = errornumber
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, etype, e, tb):
+        return (etype == WindowsError and e.winerror == self.errornumber)
+
+
+KeyValue = collections.namedtuple("KeyValue", ["name", "value", "type"])
+
+class PyHKey(object):
+    def __init__(self, surkey, name, sam=_winreg.KEY_READ):
+        self.surkey = surkey
+        self._phkey = None
+        self.name = name
+        self.fullname = self.surkey.fullname + "\\" + self.name if self.name else self.surkey.name
+        self.sam = sam
+
+    def __repr__(self):
+        return '<PyHKey "{0}">'.format(self.fullname)
+
+    @property
+    def phkey(self):
+        if self._phkey is not None:
+            return self._phkey
+        print("OPEN <{0}, {1}>".format(self.surkey.phkey, self.name))
+        self._phkey = _winreg.OpenKeyEx(self.surkey.phkey, self.name, 0, self.sam)
+        return self._phkey
+
+    @property
+    def subkeys(self):
+        res = []
+        with ExpectWindowsError(259):
+            for i in itertools.count():
+                res.append(_winreg.EnumKey(self.phkey, i))
+        return [PyHKey(self, n) for n in  res]
+
+    @property
+    def values(self):
+        res = []
+        with ExpectWindowsError(259):
+            for i in itertools.count():
+                res.append(_winreg.EnumValue(self.phkey, i))
+        return [KeyValue(*r) for r in res]
+
+    def open_subkey(self, name):
+        return PyHKey(self, name, self.sam)
+        
+    def reopen(self, new_sam):
+        return PyHKey(self.surkey, self.name, new_sam)
+        
+    __getitem__ = open_subkey
+
+class DummyPHKEY(object):
+    def __init__(self, phkey, name):
+        self.phkey = phkey
+        self.name = name
+
+
+HKEY_LOCAL_MACHINE = PyHKey(DummyPHKEY(_winreg.HKEY_LOCAL_MACHINE, "HKEY_LOCAL_MACHINE"), "", _winreg.KEY_READ)
+
+HKEY_CLASSES_ROOT = PyHKey(DummyPHKEY(_winreg.HKEY_CLASSES_ROOT, "HKEY_CLASSES_ROOT"), "", _winreg.KEY_READ )
+
+HKEY_CURRENT_USER = PyHKey(DummyPHKEY(_winreg.HKEY_CURRENT_USER, "HKEY_CURRENT_USER"), "", _winreg.KEY_READ)
+
+HKEY_DYN_DATA = PyHKey(DummyPHKEY(_winreg.HKEY_DYN_DATA, "HKEY_DYN_DATA"), "", _winreg.KEY_READ)
+
+HKEY_PERFORMANCE_DATA = PyHKey(DummyPHKEY(_winreg.HKEY_PERFORMANCE_DATA, "HKEY_PERFORMANCE_DATA"), "", _winreg.KEY_READ)
+
+HKEY_USERS = PyHKey(DummyPHKEY(_winreg.HKEY_USERS, "HKEY_USERS"), "", _winreg.KEY_READ )
+
+
+HKEY_CURRENT_USER[r"Software\Microsoft\Windows\CurrentVersion\Run"].values
diff --git a/wintrust.py b/wintrust.py
new file mode 100644
index 0000000..3c5b529
--- /dev/null
+++ b/wintrust.py
@@ -0,0 +1,69 @@
+import ctypes
+import struct
+import windows
+from windows.generated_def.winstructs import *
+
+IID_PACK = "<I", "<H", "<H", "<B", "<B", "<B", "<B", "<B", "<B", "<B", "<B"
+def get_IID_from_raw(raw):
+    s = "".join([struct.pack(i, j) for i, j in zip(IID_PACK, raw)])
+    return ctypes.create_string_buffer(s)
+
+
+WINTRUST_ACTION_GENERIC_VERIFY_V2_RAW = (0xaac56b, 0xcd44,  0x11d0,
+                    0x8c, 0xc2, 0x0, 0xc0, 0x4f, 0xc2, 0x95, 0xee)
+
+WINTRUST_ACTION_GENERIC_VERIFY_V2 = get_IID_from_raw(WINTRUST_ACTION_GENERIC_VERIFY_V2_RAW)
+
+WTD_UI_ALL    = 1
+WTD_UI_NONE   = 2
+WTD_UI_NOBAD  = 3
+WTD_UI_NOGOOD = 4
+
+WTD_REVOKE_NONE         = 0x00000000
+WTD_REVOKE_WHOLECHAIN   = 0x00000001
+
+WTD_CHOICE_FILE    = 1
+WTD_CHOICE_CATALOG = 2
+WTD_CHOICE_BLOB    = 3
+WTD_CHOICE_SIGNER  = 4
+WTD_CHOICE_CERT    = 5
+
+WTD_STATEACTION_IGNORE           = 0x00000000
+WTD_STATEACTION_VERIFY           = 0x00000001
+WTD_STATEACTION_CLOSE            = 0x00000002
+WTD_STATEACTION_AUTO_CACHE       = 0x00000003
+WTD_STATEACTION_AUTO_CACHE_FLUSH = 0x00000004
+
+def check_signature(filename):
+    print("Filename is <{0}>".format(repr(filename)))
+    file_data = WINTRUST_FILE_INFO()
+    file_data.cbStruct = ctypes.sizeof(WINTRUST_FILE_INFO)
+    file_data.pcwszFilePath = filename
+    file_data.hFile = None
+    file_data.pgKnownSubject = None
+
+    WVTPolicyGUID =  WINTRUST_ACTION_GENERIC_VERIFY_V2
+
+    win_trust_data = WINTRUST_DATA()
+
+    win_trust_data.cbStruct = ctypes.sizeof(WINTRUST_DATA)
+    win_trust_data.pPolicyCallbackData = None
+    win_trust_data.pSIPClientData = None
+    win_trust_data.dwUIChoice = WTD_UI_NONE
+    win_trust_data.fdwRevocationChecks = WTD_REVOKE_NONE
+    win_trust_data.dwUnionChoice = WTD_CHOICE_FILE
+    win_trust_data.dwStateAction = WTD_STATEACTION_VERIFY
+    win_trust_data.hWVTStateData = None
+    win_trust_data.pwszURLReference = None
+    win_trust_data.dwUIContext = 0
+    win_trust_data.tmp_union.pFile = ctypes.pointer(file_data)
+
+    WinVerifyTrust = ctypes.WinDLL("wintrust").WinVerifyTrust
+
+    x = WinVerifyTrust(None, ctypes.byref(WVTPolicyGUID), ctypes.byref(win_trust_data))
+
+    win_trust_data.dwStateAction = WTD_STATEACTION_CLOSE
+
+    WinVerifyTrust(None, ctypes.byref(WVTPolicyGUID), ctypes.byref(win_trust_data))
+
+    return x & 0xffffffff
\ No newline at end of file
diff --git a/wmi.py b/wmi.py
new file mode 100644
index 0000000..166fd2d
--- /dev/null
+++ b/wmi.py
@@ -0,0 +1,207 @@
+import windows
+import ctypes
+from ctypes.wintypes import *
+from windows.generated_def.winstructs import *
+import struct
+import functools
+
+
+# Move simple_com from LKD to windows ?
+
+IID_PACK = "<I", "<H", "<H", "<B", "<B", "<B", "<B", "<B", "<B", "<B", "<B"
+
+
+def get_IID_from_raw(raw):
+    return "".join([struct.pack(i, j) for i, j in zip(IID_PACK, raw)])
+
+
+class COMInterface(ctypes.c_void_p):
+    _functions_ = {
+        "QueryInterface": ctypes.WINFUNCTYPE(HRESULT, ctypes.c_void_p, ctypes.c_void_p)(0, "QueryInterface"),
+        "AddRef": ctypes.WINFUNCTYPE(HRESULT)(1, "AddRef"),
+        "Release": ctypes.WINFUNCTYPE(HRESULT)(2, "Release")
+    }
+
+    def __getattr__(self, name):
+        if name in self._functions_:
+            return functools.partial(self._functions_[name], self)
+        return super(COMInterface, self).__getattribute__(name)
+
+
+class IWbemLocator(COMInterface):
+    _functions_ = {
+        "ConnectServer": ctypes.WINFUNCTYPE(HRESULT, ctypes.c_wchar_p, ctypes.c_wchar_p, ctypes.c_wchar_p, ctypes.c_wchar_p,
+                                                ctypes.c_long, ctypes.c_wchar_p, ctypes.c_void_p, POINTER(ctypes.c_void_p))(3, "ConnectServer")
+    }
+
+
+class IWbemServices(COMInterface):
+    _functions_ = {
+        "ExecQuery": ctypes.WINFUNCTYPE(HRESULT, ctypes.c_wchar_p, ctypes.c_wchar_p, ctypes.c_long, ctypes.c_void_p, POINTER(ctypes.c_void_p))(20, "ExecQuery")
+    }
+
+
+class IEnumWbemClassObject(COMInterface):
+    _functions_ = {
+        "Next": ctypes.WINFUNCTYPE(HRESULT, ctypes.c_long, ctypes.c_ulong, POINTER(ctypes.c_void_p), POINTER(ctypes.c_long))(4, "Next")
+    }
+
+
+class IWbemClassObject(COMInterface):
+    _functions_ = {
+        "Get": ctypes.WINFUNCTYPE(HRESULT, ctypes.c_wchar_p, ctypes.c_long, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_long)(4, "Get"),
+        "GetNames": ctypes.WINFUNCTYPE(HRESULT, ctypes.c_wchar_p, ctypes.c_long, ctypes.c_void_p, POINTER(ctypes.c_void_p))(7, "GetNames")
+    }
+
+
+CLSID_WbemAdministrativeLocator_raw = 0xcb8555cc, 0x9128, 0x11d1, 0xad, 0x9b, 0x00, 0xc0, 0x4f, 0xd8, 0x0fd, 0xff
+CLSID_WbemAdministrativeLocator_IID = get_IID_from_raw(CLSID_WbemAdministrativeLocator_raw)
+
+IID_IWbemLocator_raw = 0x0DC12A687, 0x737F, 0x11CF, 0x88, 0x4d, 0x00, 0xaa, 0x00, 0x4b, 0x2e, 0x24
+IID_IWbemLocator= get_IID_from_raw(IID_IWbemLocator_raw)
+
+
+BSTR = ctypes.c_wchar_p
+
+
+class _tagBRECORD(ctypes.Structure):
+    _fields_ = [("pvRecord", PVOID), ("pRecInfo", PVOID)]
+
+    
+class SAFEARRAY(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
+
+
+class SimpleVariantData(ctypes.Union):
+    _fields_ = [("llVal", LONGLONG),
+                ("lVal", LONG),
+                ("bVal", BYTE),
+                ("iVal", SHORT),
+                ("fltVal", FLOAT),
+                ("dblVal", DOUBLE),
+                ("bstrVal", BSTR),
+                ("pbstrVal", POINTER(BSTR)),
+                ("byref", PVOID),
+                ("parray", POINTER(SAFEARRAY)),
+                ("pbyref", PVOID),
+                ("ullVal", ULONGLONG),
+                ("llVal", LONGLONG),
+                ("__VARIANT_NAME_4", _tagBRECORD)]
+
+                
+# TODO: put real struct in winstruct (need a real parser for generation)
+class SimpleVariant(ctypes.Structure):
+    _fields_ = [("vt", WORD), ("wReserved1", WORD), ("wReserved2", WORD), ("wReserved3", WORD),  ("_Data", SimpleVariantData)]
+
+    @property
+    def asbstr(self):
+        if self.vt != VT_BSTR:
+            raise ValueError("asbstr on non-bstr variant")
+        return self._Data.bstrVal
+
+    @property
+    def aslong(self):
+        if not self.vt in [VT_I4, VT_BOOL]:
+            raise ValueError("aslong on non-long variant")
+        return self._Data.lVal
+
+    @property
+    def asbool(self):
+        if not self.vt in [VT_BOOL]:
+            raise ValueError("get_bstr on non-bool variant")
+        return bool(self.aslong)
+
+
+class WmiRequester(object):
+    def __init__(self):
+        locator = IWbemLocator()
+        service = IWbemServices()
+
+        assert ctypes.windll.ole32.CoInitializeEx(0, 0) == 0
+        assert ctypes.windll.ole32.CoInitializeSecurity(0, -1, 0,0, 0, 3, 0,0,0) == 0
+        assert ctypes.windll.ole32.CoCreateInstance(CLSID_WbemAdministrativeLocator_IID, 0, 1, IID_IWbemLocator, ctypes.byref(locator)) == 0
+
+        locator.ConnectServer("root\\cimv2", None, None , None, 0x80, None, None, ctypes.byref(service))
+        self.service = service
+
+
+    def request_select(self, frm, attrs):
+        enumerator = IEnumWbemClassObject()
+        self.service.ExecQuery("WQL", "select * from {0}".format(frm), 0x20, 0, ctypes.byref(enumerator))
+
+        count = ctypes.c_long(0)
+        processor = IWbemClassObject()
+        res = []
+        enumerator.Next(0xffffffff, 1, ctypes.byref(processor), ctypes.byref(count))
+        while count.value:
+            current_res = {}
+            variant_res = SimpleVariant()
+            self.get_names(processor)
+            if attrs == "*":
+                attrs = [x for x in self.get_names(processor) if not x.startswith("__")]
+            for name in attrs:
+                processor.Get(name, 0, ctypes.byref(variant_res), 0, 0)
+                # 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)
+                    if variant_res.vt & VT_TYPEMASK == VT_I4:
+                        current_res[name] = variant_res._Data.parray[0].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
+                else:
+                    print("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
+
+    def get_names(self, processor):
+        res = PVOID()
+        processor.GetNames(None, 0, None, res)
+        return SAFEARRAY.of_type(res.value, BSTR).to_list()
+
+
+req = WmiRequester()
+
+#v = req.request_select("Win32_Process", ["Name", "CommandLine", "ExecutablePath"])
+#import pprint
+#
+#pprint.pprint(v)
+#
+#
+#v = req.request_select("Win32_StartupCommand", ["Command"]);
+#print("========")
+#pprint.pprint(v)
+#
+#v = req.request_select("Win32_ComputerSystemProduct", ["Name"]);
+#print("========")
+#pprint.pprint(v)
+
+v = req.request_select("Win32_Bios", "*")