pythongh-89545: Adds internal _wmi module on Windows for directly que…

…rying OS properties

Misc/NEWS.d/next/Windows/2022-08-26-00-11-18.gh-issue-89545.zmJMY_.rst

@@ -0,0 +1 @@
+Adds internal ``_wmi`` module for directly querying OS properties

PC/_wmimodule.cpp

@@ -0,0 +1,277 @@
+//
+// Helper library for querying WMI using its COM-based query API.
+//
+// Copyright (c) Microsoft Corporation
+// Licensed to PSF under a contributor agreement
+//
+
+// Version history
+//  2022-08: Initial contribution (Steve Dower)
+
+#define _WIN32_DCOM
+#include <Windows.h>
+#include <comdef.h>
+#include <Wbemidl.h>
+#include <propvarutil.h>
+
+#include <Python.h>
+#include "clinic/_wmimodule.cpp.h"
+
+
+/*[clinic input]
+module _wmi
+[clinic start generated code]*/
+/*[clinic end generated code: output=da39a3ee5e6b4b0d input=7ca95dad1453d10d]*/
+
+
+
+struct _query_data {
+    LPCWSTR query;
+    HANDLE writePipe;
+    HANDLE readPipe;
+};
+
+
+static DWORD WINAPI
+_query_thread(LPVOID param)
+{
+    IWbemLocator *locator = NULL;
+    IWbemServices *services = NULL;
+    IEnumWbemClassObject* enumerator = NULL;
+    BSTR bstrQuery = NULL;
+    struct _query_data *data = (struct _query_data*)param;
+
+    HRESULT hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
+    if (FAILED(hr)) {
+        CloseHandle(data->writePipe);
+        return (DWORD)hr;
+    }
+
+    hr = CoInitializeSecurity(
+        NULL, -1, NULL, NULL,
+        RPC_C_AUTHN_LEVEL_DEFAULT,
+        RPC_C_IMP_LEVEL_IMPERSONATE,
+        NULL, EOAC_NONE, NULL
+    );
+    if (SUCCEEDED(hr)) {
+        hr = CoCreateInstance(
+            CLSID_WbemLocator, 0, CLSCTX_INPROC_SERVER,
+            IID_IWbemLocator, (LPVOID *)&locator
+        );
+    }
+    if (SUCCEEDED(hr)) {
+        hr = locator->ConnectServer(
+            bstr_t(L"ROOT\\CIMV2"),
+            NULL, NULL, 0, NULL, 0, 0, &services
+        );
+    }
+    if (SUCCEEDED(hr)) {
+        hr = CoSetProxyBlanket(
+            services, RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE, NULL,
+            RPC_C_AUTHN_LEVEL_CALL, RPC_C_IMP_LEVEL_IMPERSONATE,
+            NULL, EOAC_NONE
+        );
+    }
+    if (SUCCEEDED(hr)) {
+        bstrQuery = SysAllocString(data->query);
+        if (!bstrQuery) {
+            hr = HRESULT_FROM_WIN32(ERROR_NOT_ENOUGH_MEMORY);
+        }
+    }
+    if (SUCCEEDED(hr)) {
+        hr = services->ExecQuery(
+            bstr_t("WQL"), 
+            bstrQuery,
+            WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, 
+            NULL,
+            &enumerator
+        );
+    }
+
+    // Okay, after all that, at this stage we should have an enumerator
+    // to the query results and can start writing them to the pipe!
+    IWbemClassObject *value = NULL;
+    int endOfEnum = FALSE;
+    while (SUCCEEDED(hr) && !endOfEnum) {
+        ULONG got = 0;
+        hr = enumerator->Next(WBEM_INFINITE, 1, &value, &got);
+        if (hr == WBEM_S_FALSE) {
+            // Could be at the end, but still got a result this time
+            endOfEnum = TRUE;
+            hr = 0;
+            break;
+        }
+        if (FAILED(hr) || got != 1 || !value) {
+            continue;
+        }
+        // Okay, now we have each resulting object it's time to
+        // enumerate its members
+        hr = value->BeginEnumeration(0);
+        while (SUCCEEDED(hr)) {
+            BSTR propName;
+            VARIANT propValue;
+            long flavor;
+            hr = value->Next(0, &propName, &propValue, NULL, &flavor);
+            if (hr == WBEM_S_NO_MORE_DATA) {
+                hr = 0;
+                break;
+            }
+            if (SUCCEEDED(hr) && (flavor & WBEM_FLAVOR_MASK_ORIGIN) != WBEM_FLAVOR_ORIGIN_SYSTEM) {
+                WCHAR propStr[8192];
+                hr = VariantToString(propValue, propStr, sizeof(propStr) / sizeof(propStr[0]));
+                if (SUCCEEDED(hr)) {
+                    DWORD cbStr;
+                    DWORD written;
+                    cbStr = (DWORD)(wcslen(propName) * sizeof(propName[0]));
+                    WriteFile(data->writePipe, propName, cbStr, &written, NULL);
+                    WriteFile(data->writePipe, (LPVOID)L"=", 2, &written, NULL);
+                    cbStr = (DWORD)(wcslen(propStr) * sizeof(propStr[0]));
+                    WriteFile(data->writePipe, propStr, cbStr, &written, NULL);
+                    WriteFile(data->writePipe, (LPVOID)L"\0", 2, &written, NULL);
+                }
+                VariantClear(&propValue);
+                SysFreeString(propName);
+            }
+        }
+        value->EndEnumeration();
+        value->Release();
+    }
+
+    if (bstrQuery) {
+        SysFreeString(bstrQuery);
+    }
+    if (enumerator) {
+        enumerator->Release();
+    }
+    if (services) {
+        services->Release();
+    }
+    if (locator) {
+        locator->Release();
+    }
+    CoUninitialize();
+    CloseHandle(data->writePipe);
+    return (DWORD)hr;
+}
+
+
+/*[clinic input]
+_wmi.exec_query
+
+    query: unicode
+
+Runs a WMI query against the local machine.
+
+This returns a single string with 'name=value' pairs in a flat array separated
+by null characters.
+[clinic start generated code]*/
+
+static PyObject *
+_wmi_exec_query_impl(PyObject *module, PyObject *query)
+/*[clinic end generated code: output=a62303d5bb5e003f input=48d2d0a1e1a7e3c2]*/
+
+/*[clinic end generated code]*/
+{
+    PyObject *result = NULL;
+    HANDLE hThread = NULL;
+    int err = 0;
+    WCHAR buffer[16384];
+    DWORD offset = 0;
+    DWORD bytesRead;
+    struct _query_data data = {0};
+
+    if (PySys_Audit("_wmi.exec_query", "O", query) < 0) {
+        return NULL;
+    }
+
+    data.query = PyUnicode_AsWideCharString(query, NULL);
+    if (!data.query) {
+        return NULL;
+    }
+
+    Py_BEGIN_ALLOW_THREADS
+
+    if (!CreatePipe(&data.readPipe, &data.writePipe, NULL, 0) ||
+        !(hThread = CreateThread(NULL, 0, _query_thread, (LPVOID*)&data, 0, NULL))) {
+        err = GetLastError();
+    }
+
+    while (!err) {
+        if (ReadFile(
+            data.readPipe,
+            (LPVOID)&buffer[offset / sizeof(buffer[0])],
+            sizeof(buffer) - offset,
+            &bytesRead,
+            NULL
+        )) {
+            offset += bytesRead;
+            if (offset >= sizeof(buffer)) {
+                err = ERROR_MORE_DATA;
+            }
+        } else {
+            err = GetLastError();
+        }
+    }
+
+    if (data.readPipe) {
+        CloseHandle(data.readPipe);
+    }
+    if (data.writePipe) {
+        CloseHandle(data.writePipe);
+    }
+
+    if (err == ERROR_BROKEN_PIPE) {
+        // broken pipe indicates some kind of failure, but the real error
+        // code will come as the thread exit
+        if (WaitForSingleObject(hThread, 1000) != WAIT_OBJECT_0 ||
+            !GetExitCodeThread(hThread, (LPDWORD)&err)) {
+            err = GetLastError();
+        }
+    }
+
+    CloseHandle(hThread);
+    hThread = NULL;
+
+    Py_END_ALLOW_THREADS
+
+    PyMem_Free((void *)data.query);
+
+    return PyUnicode_FromWideChar(buffer, offset  / sizeof(buffer[0]) - 1);
+}
+
+
+static PyMethodDef wmi_functions[] = {
+    _WMI_EXEC_QUERY_METHODDEF
+    { NULL, NULL, 0, NULL }
+};
+
+static int exec_wmi(PyObject *module)
+{
+    PyModule_AddFunctions(module, wmi_functions);
+
+    return 0; // success
+}
+
+static PyModuleDef_Slot wmi_slots[] = {
+    { Py_mod_exec, exec_wmi },
+    { 0, NULL }
+};
+
+static PyModuleDef wmi_def = {
+    PyModuleDef_HEAD_INIT,
+    "_wmi",
+    NULL, // doc
+    0,              // m_size
+    NULL,           // m_methods
+    wmi_slots,
+    NULL,           // m_traverse
+    NULL,           // m_clear
+    NULL,           // m_free
+};
+
+extern "C" {
+    PyMODINIT_FUNC PyInit__wmi(void)
+    {
+        return PyModuleDef_Init(&wmi_def);
+    }
+}

PC/clinic/_wmimodule.cpp.h

@@ -0,0 +1,75 @@
+/*[clinic input]
+preserve
+[clinic start generated code]*/
+
+#if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE)
+#  include "pycore_gc.h"            // PyGC_Head
+#  include "pycore_runtime.h"       // _Py_ID()
+#endif
+
+
+PyDoc_STRVAR(_wmi_exec_query__doc__,
+"exec_query($module, /, query)\n"
+"--\n"
+"\n"
+"Runs a WMI query against the local machine.\n"
+"\n"
+"This returns a single string with \'name=value\' pairs in a flat array separated\n"
+"by null characters.");
+
+#define _WMI_EXEC_QUERY_METHODDEF    \
+    {"exec_query", _PyCFunction_CAST(_wmi_exec_query), METH_FASTCALL|METH_KEYWORDS, _wmi_exec_query__doc__},
+
+static PyObject *
+_wmi_exec_query_impl(PyObject *module, PyObject *query);
+
+static PyObject *
+_wmi_exec_query(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    PyObject *return_value = NULL;
+    #if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE)
+
+    #define NUM_KEYWORDS 1
+    static struct {
+        PyGC_Head _this_is_not_used;
+        PyObject_VAR_HEAD
+        PyObject *ob_item[NUM_KEYWORDS];
+    } _kwtuple = {
+        .ob_base = PyVarObject_HEAD_INIT(&PyTuple_Type, NUM_KEYWORDS)
+        .ob_item = { &_Py_ID(query), },
+    };
+    #undef NUM_KEYWORDS
+    #define KWTUPLE (&_kwtuple.ob_base.ob_base)
+
+    #else  // !Py_BUILD_CORE
+    #  define KWTUPLE NULL
+    #endif  // !Py_BUILD_CORE
+
+    static const char * const _keywords[] = {"query", NULL};
+    static _PyArg_Parser _parser = {
+        .keywords = _keywords,
+        .fname = "exec_query",
+        .kwtuple = KWTUPLE,
+    };
+    #undef KWTUPLE
+    PyObject *argsbuf[1];
+    PyObject *query;
+
+    args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 1, 1, 0, argsbuf);
+    if (!args) {
+        goto exit;
+    }
+    if (!PyUnicode_Check(args[0])) {
+        _PyArg_BadArgument("exec_query", "argument 'query'", "str", args[0]);
+        goto exit;
+    }
+    if (PyUnicode_READY(args[0]) == -1) {
+        goto exit;
+    }
+    query = args[0];
+    return_value = _wmi_exec_query_impl(module, query);
+
+exit:
+    return return_value;
+}
+/*[clinic end generated code: output=7fdf0c0579ddb566 input=a9049054013a1b77]*/