| // | |
| // 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) | |
| // clinic/_wmimodule.cpp.h uses internal pycore_modsupport.h API | |
| #ifndef Py_BUILD_CORE_BUILTIN | |
| # define Py_BUILD_CORE_MODULE 1 | |
| #endif | |
| #define _WIN32_DCOM | |
| #include <Windows.h> | |
| #include <comdef.h> | |
| #include <Wbemidl.h> | |
| #include <propvarutil.h> | |
| #include <Python.h> | |
| #if _MSVC_LANG >= 202002L | |
| // We can use clinic directly when the C++ compiler supports C++20 | |
| #include "clinic/_wmimodule.cpp.h" | |
| #else | |
| // Cannot use clinic because of missing C++20 support, so create a simpler | |
| // API instead. This won't impact releases, so fine to omit the docstring. | |
| static PyObject *_wmi_exec_query_impl(PyObject *module, PyObject *query); | |
| #define _WMI_EXEC_QUERY_METHODDEF {"exec_query", _wmi_exec_query_impl, METH_O, NULL}, | |
| #endif | |
| /*[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; | |
| HANDLE initEvent; | |
| HANDLE connectEvent; | |
| }; | |
| static DWORD WINAPI | |
| _query_thread(LPVOID param) | |
| { | |
| IWbemLocator *locator = NULL; | |
| IWbemServices *services = NULL; | |
| IEnumWbemClassObject* enumerator = NULL; | |
| HRESULT hr = S_OK; | |
| BSTR bstrQuery = NULL; | |
| _query_data data = *(struct _query_data*)param; | |
| // gh-125315: Copy the query string first, so that if the main thread gives | |
| // up on waiting we aren't left with a dangling pointer (and a likely crash) | |
| bstrQuery = SysAllocString(data.query); | |
| if (!bstrQuery) { | |
| hr = HRESULT_FROM_WIN32(ERROR_NOT_ENOUGH_MEMORY); | |
| } | |
| if (SUCCEEDED(hr)) { | |
| hr = CoInitializeEx(nullptr, COINIT_APARTMENTTHREADED); | |
| } | |
| if (FAILED(hr)) { | |
| CloseHandle(data.writePipe); | |
| if (bstrQuery) { | |
| SysFreeString(bstrQuery); | |
| } | |
| return (DWORD)hr; | |
| } | |
| hr = CoInitializeSecurity( | |
| NULL, -1, NULL, NULL, | |
| RPC_C_AUTHN_LEVEL_DEFAULT, | |
| RPC_C_IMP_LEVEL_IMPERSONATE, | |
| NULL, EOAC_NONE, NULL | |
| ); | |
| // gh-96684: CoInitializeSecurity will fail if another part of the app has | |
| // already called it. Hopefully they passed lenient enough settings that we | |
| // can complete the WMI query, so keep going. | |
| if (hr == RPC_E_TOO_LATE) { | |
| hr = 0; | |
| } | |
| if (SUCCEEDED(hr)) { | |
| hr = CoCreateInstance( | |
| CLSID_WbemLocator, 0, CLSCTX_INPROC_SERVER, | |
| IID_IWbemLocator, (LPVOID *)&locator | |
| ); | |
| } | |
| if (SUCCEEDED(hr) && !SetEvent(data.initEvent)) { | |
| hr = HRESULT_FROM_WIN32(GetLastError()); | |
| } | |
| if (SUCCEEDED(hr)) { | |
| hr = locator->ConnectServer( | |
| bstr_t(L"ROOT\\CIMV2"), | |
| NULL, NULL, 0, NULL, 0, 0, &services | |
| ); | |
| } | |
| if (SUCCEEDED(hr) && !SetEvent(data.connectEvent)) { | |
| hr = HRESULT_FROM_WIN32(GetLastError()); | |
| } | |
| 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)) { | |
| 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 startOfEnum = TRUE; | |
| int endOfEnum = FALSE; | |
| while (SUCCEEDED(hr) && !endOfEnum) { | |
| ULONG got = 0; | |
| DWORD written; | |
| 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; | |
| } | |
| if (!startOfEnum && !WriteFile(data.writePipe, (LPVOID)L"\0", 2, &written, NULL)) { | |
| hr = HRESULT_FROM_WIN32(GetLastError()); | |
| break; | |
| } | |
| startOfEnum = FALSE; | |
| // Okay, now we have each resulting object it's time to | |
| // enumerate its members | |
| hr = value->BeginEnumeration(0); | |
| if (FAILED(hr)) { | |
| value->Release(); | |
| break; | |
| } | |
| 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 cbStr1, cbStr2; | |
| cbStr1 = (DWORD)(wcslen(propName) * sizeof(propName[0])); | |
| cbStr2 = (DWORD)(wcslen(propStr) * sizeof(propStr[0])); | |
| if (!WriteFile(data.writePipe, propName, cbStr1, &written, NULL) || | |
| !WriteFile(data.writePipe, (LPVOID)L"=", 2, &written, NULL) || | |
| !WriteFile(data.writePipe, propStr, cbStr2, &written, NULL) || | |
| !WriteFile(data.writePipe, (LPVOID)L"\0", 2, &written, NULL) | |
| ) { | |
| hr = HRESULT_FROM_WIN32(GetLastError()); | |
| } | |
| } | |
| 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; | |
| } | |
| static DWORD | |
| wait_event(HANDLE event, DWORD timeout) | |
| { | |
| DWORD err = 0; | |
| switch (WaitForSingleObject(event, timeout)) { | |
| case WAIT_OBJECT_0: | |
| break; | |
| case WAIT_TIMEOUT: | |
| err = WAIT_TIMEOUT; | |
| break; | |
| default: | |
| err = GetLastError(); | |
| break; | |
| } | |
| return err; | |
| } | |
| /*[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=a8d5710acdfbf515]*/ | |
| /*[clinic end generated code]*/ | |
| { | |
| HANDLE hThread = NULL; | |
| int err = 0; | |
| WCHAR buffer[8192]; | |
| 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; | |
| } | |
| if (0 != _wcsnicmp(data.query, L"select ", 7)) { | |
| PyMem_Free((void *)data.query); | |
| PyErr_SetString(PyExc_ValueError, "only SELECT queries are supported"); | |
| return NULL; | |
| } | |
| Py_BEGIN_ALLOW_THREADS | |
| data.initEvent = CreateEvent(NULL, TRUE, FALSE, NULL); | |
| data.connectEvent = CreateEvent(NULL, TRUE, FALSE, NULL); | |
| if (!data.initEvent || !data.connectEvent || | |
| !CreatePipe(&data.readPipe, &data.writePipe, NULL, 0)) | |
| { | |
| err = GetLastError(); | |
| } else { | |
| hThread = CreateThread(NULL, 0, _query_thread, (LPVOID*)&data, 0, NULL); | |
| if (!hThread) { | |
| err = GetLastError(); | |
| // Normally the thread proc closes this handle, but since we never started | |
| // we need to close it here. | |
| CloseHandle(data.writePipe); | |
| } | |
| } | |
| // gh-112278: If current user doesn't have permission to query the WMI, the | |
| // function IWbemLocator::ConnectServer will hang for 5 seconds, and there | |
| // is no way to specify the timeout. So we use an Event object to simulate | |
| // a timeout. The initEvent will be set after COM initialization, it will | |
| // take a longer time when first initialized. The connectEvent will be set | |
| // after connected to WMI. | |
| if (!err) { | |
| err = wait_event(data.initEvent, 1000); | |
| if (!err) { | |
| err = wait_event(data.connectEvent, 100); | |
| } | |
| } | |
| 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 (hThread) { | |
| // Allow the thread some time to clean up | |
| int thread_err; | |
| switch (WaitForSingleObject(hThread, 100)) { | |
| case WAIT_OBJECT_0: | |
| // Thread ended cleanly | |
| if (!GetExitCodeThread(hThread, (LPDWORD)&thread_err)) { | |
| thread_err = GetLastError(); | |
| } | |
| break; | |
| case WAIT_TIMEOUT: | |
| // Probably stuck - there's not much we can do, unfortunately | |
| thread_err = WAIT_TIMEOUT; | |
| break; | |
| default: | |
| thread_err = GetLastError(); | |
| break; | |
| } | |
| // An error on our side is more likely to be relevant than one from | |
| // the thread, but if we don't have one on our side we'll take theirs. | |
| if (err == 0 || err == ERROR_BROKEN_PIPE) { | |
| err = thread_err; | |
| } | |
| CloseHandle(hThread); | |
| } | |
| CloseHandle(data.initEvent); | |
| CloseHandle(data.connectEvent); | |
| hThread = NULL; | |
| Py_END_ALLOW_THREADS | |
| PyMem_Free((void *)data.query); | |
| if (err == ERROR_MORE_DATA) { | |
| PyErr_Format(PyExc_OSError, "Query returns more than %zd characters", Py_ARRAY_LENGTH(buffer)); | |
| return NULL; | |
| } else if (err) { | |
| PyErr_SetFromWindowsErr(err); | |
| return NULL; | |
| } | |
| if (!offset) { | |
| return PyUnicode_FromStringAndSize(NULL, 0); | |
| } | |
| return PyUnicode_FromWideChar(buffer, offset / sizeof(buffer[0]) - 1); | |
| } | |
| static PyMethodDef wmi_functions[] = { | |
| _WMI_EXEC_QUERY_METHODDEF | |
| { NULL, NULL, 0, NULL } | |
| }; | |
| static PyModuleDef_Slot wmi_slots[] = { | |
| {Py_mod_gil, Py_MOD_GIL_NOT_USED}, | |
| {0, NULL}, | |
| }; | |
| static PyModuleDef wmi_def = { | |
| PyModuleDef_HEAD_INIT, | |
| "_wmi", | |
| NULL, // doc | |
| 0, // m_size | |
| wmi_functions, // m_methods | |
| wmi_slots, // m_slots | |
| }; | |
| extern "C" { | |
| PyMODINIT_FUNC PyInit__wmi(void) | |
| { | |
| return PyModuleDef_Init(&wmi_def); | |
| } | |
| } |