1 //
2 // Helper library for querying WMI using its COM-based query API.
3 //
4 // Copyright (c) Microsoft Corporation
5 // Licensed to PSF under a contributor agreement
6 //
7
8 // Version history
9 // 2022-08: Initial contribution (Steve Dower)
10
11 // clinic/_wmimodule.cpp.h uses internal pycore_modsupport.h API
12 #ifndef Py_BUILD_CORE_BUILTIN
13 # define Py_BUILD_CORE_MODULE 1
14 #endif
15
16 #define _WIN32_DCOM
17 #include <Windows.h>
18 #include <comdef.h>
19 #include <Wbemidl.h>
20 #include <propvarutil.h>
21
22 #include <Python.h>
23
24
25 #if _MSVC_LANG >= 202002L
26 // We can use clinic directly when the C++ compiler supports C++20
27 #include "clinic/_wmimodule.cpp.h"
28 #else
29 // Cannot use clinic because of missing C++20 support, so create a simpler
30 // API instead. This won't impact releases, so fine to omit the docstring.
31 static PyObject *_wmi_exec_query_impl(PyObject *module, PyObject *query);
32 #define _WMI_EXEC_QUERY_METHODDEF {"exec_query", _wmi_exec_query_impl, METH_O, NULL},
33 #endif
34
35
36 /*[clinic input]
37 module _wmi
38 [clinic start generated code]*/
39 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=7ca95dad1453d10d]*/
40
41
42
43 struct _query_data {
44 LPCWSTR query;
45 HANDLE writePipe;
46 HANDLE readPipe;
47 HANDLE initEvent;
48 HANDLE connectEvent;
49 };
50
51
52 static DWORD WINAPI
_query_thread(LPVOID param)53 _query_thread(LPVOID param)
54 {
55 IWbemLocator *locator = NULL;
56 IWbemServices *services = NULL;
57 IEnumWbemClassObject* enumerator = NULL;
58 HRESULT hr = S_OK;
59 BSTR bstrQuery = NULL;
60 struct _query_data *data = (struct _query_data*)param;
61
62 // gh-125315: Copy the query string first, so that if the main thread gives
63 // up on waiting we aren't left with a dangling pointer (and a likely crash)
64 bstrQuery = SysAllocString(data->query);
65 if (!bstrQuery) {
66 hr = HRESULT_FROM_WIN32(ERROR_NOT_ENOUGH_MEMORY);
67 }
68
69 if (SUCCEEDED(hr)) {
70 hr = CoInitializeEx(nullptr, COINIT_APARTMENTTHREADED);
71 }
72
73 if (FAILED(hr)) {
74 CloseHandle(data->writePipe);
75 if (bstrQuery) {
76 SysFreeString(bstrQuery);
77 }
78 return (DWORD)hr;
79 }
80
81 hr = CoInitializeSecurity(
82 NULL, -1, NULL, NULL,
83 RPC_C_AUTHN_LEVEL_DEFAULT,
84 RPC_C_IMP_LEVEL_IMPERSONATE,
85 NULL, EOAC_NONE, NULL
86 );
87 // gh-96684: CoInitializeSecurity will fail if another part of the app has
88 // already called it. Hopefully they passed lenient enough settings that we
89 // can complete the WMI query, so keep going.
90 if (hr == RPC_E_TOO_LATE) {
91 hr = 0;
92 }
93 if (SUCCEEDED(hr)) {
94 hr = CoCreateInstance(
95 CLSID_WbemLocator, 0, CLSCTX_INPROC_SERVER,
96 IID_IWbemLocator, (LPVOID *)&locator
97 );
98 }
99 if (SUCCEEDED(hr) && !SetEvent(data->initEvent)) {
100 hr = HRESULT_FROM_WIN32(GetLastError());
101 }
102 if (SUCCEEDED(hr)) {
103 hr = locator->ConnectServer(
104 bstr_t(L"ROOT\\CIMV2"),
105 NULL, NULL, 0, NULL, 0, 0, &services
106 );
107 }
108 if (SUCCEEDED(hr) && !SetEvent(data->connectEvent)) {
109 hr = HRESULT_FROM_WIN32(GetLastError());
110 }
111 if (SUCCEEDED(hr)) {
112 hr = CoSetProxyBlanket(
113 services, RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE, NULL,
114 RPC_C_AUTHN_LEVEL_CALL, RPC_C_IMP_LEVEL_IMPERSONATE,
115 NULL, EOAC_NONE
116 );
117 }
118 if (SUCCEEDED(hr)) {
119 hr = services->ExecQuery(
120 bstr_t("WQL"),
121 bstrQuery,
122 WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
123 NULL,
124 &enumerator
125 );
126 }
127
128 // Okay, after all that, at this stage we should have an enumerator
129 // to the query results and can start writing them to the pipe!
130 IWbemClassObject *value = NULL;
131 int startOfEnum = TRUE;
132 int endOfEnum = FALSE;
133 while (SUCCEEDED(hr) && !endOfEnum) {
134 ULONG got = 0;
135 DWORD written;
136 hr = enumerator->Next(WBEM_INFINITE, 1, &value, &got);
137 if (hr == WBEM_S_FALSE) {
138 // Could be at the end, but still got a result this time
139 endOfEnum = TRUE;
140 hr = 0;
141 break;
142 }
143 if (FAILED(hr) || got != 1 || !value) {
144 continue;
145 }
146 if (!startOfEnum && !WriteFile(data->writePipe, (LPVOID)L"\0", 2, &written, NULL)) {
147 hr = HRESULT_FROM_WIN32(GetLastError());
148 break;
149 }
150 startOfEnum = FALSE;
151 // Okay, now we have each resulting object it's time to
152 // enumerate its members
153 hr = value->BeginEnumeration(0);
154 if (FAILED(hr)) {
155 value->Release();
156 break;
157 }
158 while (SUCCEEDED(hr)) {
159 BSTR propName;
160 VARIANT propValue;
161 long flavor;
162 hr = value->Next(0, &propName, &propValue, NULL, &flavor);
163 if (hr == WBEM_S_NO_MORE_DATA) {
164 hr = 0;
165 break;
166 }
167 if (SUCCEEDED(hr) && (flavor & WBEM_FLAVOR_MASK_ORIGIN) != WBEM_FLAVOR_ORIGIN_SYSTEM) {
168 WCHAR propStr[8192];
169 hr = VariantToString(propValue, propStr, sizeof(propStr) / sizeof(propStr[0]));
170 if (SUCCEEDED(hr)) {
171 DWORD cbStr1, cbStr2;
172 cbStr1 = (DWORD)(wcslen(propName) * sizeof(propName[0]));
173 cbStr2 = (DWORD)(wcslen(propStr) * sizeof(propStr[0]));
174 if (!WriteFile(data->writePipe, propName, cbStr1, &written, NULL) ||
175 !WriteFile(data->writePipe, (LPVOID)L"=", 2, &written, NULL) ||
176 !WriteFile(data->writePipe, propStr, cbStr2, &written, NULL) ||
177 !WriteFile(data->writePipe, (LPVOID)L"\0", 2, &written, NULL)
178 ) {
179 hr = HRESULT_FROM_WIN32(GetLastError());
180 }
181 }
182 VariantClear(&propValue);
183 SysFreeString(propName);
184 }
185 }
186 value->EndEnumeration();
187 value->Release();
188 }
189
190 if (bstrQuery) {
191 SysFreeString(bstrQuery);
192 }
193 if (enumerator) {
194 enumerator->Release();
195 }
196 if (services) {
197 services->Release();
198 }
199 if (locator) {
200 locator->Release();
201 }
202 CoUninitialize();
203 CloseHandle(data->writePipe);
204 return (DWORD)hr;
205 }
206
207
208 static DWORD
wait_event(HANDLE event,DWORD timeout)209 wait_event(HANDLE event, DWORD timeout)
210 {
211 DWORD err = 0;
212 switch (WaitForSingleObject(event, timeout)) {
213 case WAIT_OBJECT_0:
214 break;
215 case WAIT_TIMEOUT:
216 err = WAIT_TIMEOUT;
217 break;
218 default:
219 err = GetLastError();
220 break;
221 }
222 return err;
223 }
224
225
226 /*[clinic input]
227 _wmi.exec_query
228
229 query: unicode
230
231 Runs a WMI query against the local machine.
232
233 This returns a single string with 'name=value' pairs in a flat array separated
234 by null characters.
235 [clinic start generated code]*/
236
237 static PyObject *
_wmi_exec_query_impl(PyObject * module,PyObject * query)238 _wmi_exec_query_impl(PyObject *module, PyObject *query)
239 /*[clinic end generated code: output=a62303d5bb5e003f input=48d2d0a1e1a7e3c2]*/
240
241 /*[clinic end generated code]*/
242 {
243 PyObject *result = NULL;
244 HANDLE hThread = NULL;
245 int err = 0;
246 WCHAR buffer[8192];
247 DWORD offset = 0;
248 DWORD bytesRead;
249 struct _query_data data = {0};
250
251 if (PySys_Audit("_wmi.exec_query", "O", query) < 0) {
252 return NULL;
253 }
254
255 data.query = PyUnicode_AsWideCharString(query, NULL);
256 if (!data.query) {
257 return NULL;
258 }
259
260 if (0 != _wcsnicmp(data.query, L"select ", 7)) {
261 PyMem_Free((void *)data.query);
262 PyErr_SetString(PyExc_ValueError, "only SELECT queries are supported");
263 return NULL;
264 }
265
266 Py_BEGIN_ALLOW_THREADS
267
268 data.initEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
269 data.connectEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
270 if (!data.initEvent || !data.connectEvent ||
271 !CreatePipe(&data.readPipe, &data.writePipe, NULL, 0))
272 {
273 err = GetLastError();
274 } else {
275 hThread = CreateThread(NULL, 0, _query_thread, (LPVOID*)&data, 0, NULL);
276 if (!hThread) {
277 err = GetLastError();
278 // Normally the thread proc closes this handle, but since we never started
279 // we need to close it here.
280 CloseHandle(data.writePipe);
281 }
282 }
283
284 // gh-112278: If current user doesn't have permission to query the WMI, the
285 // function IWbemLocator::ConnectServer will hang for 5 seconds, and there
286 // is no way to specify the timeout. So we use an Event object to simulate
287 // a timeout. The initEvent will be set after COM initialization, it will
288 // take a longer time when first initialized. The connectEvent will be set
289 // after connected to WMI.
290 if (!err) {
291 err = wait_event(data.initEvent, 1000);
292 if (!err) {
293 err = wait_event(data.connectEvent, 100);
294 }
295 }
296
297 while (!err) {
298 if (ReadFile(
299 data.readPipe,
300 (LPVOID)&buffer[offset / sizeof(buffer[0])],
301 sizeof(buffer) - offset,
302 &bytesRead,
303 NULL
304 )) {
305 offset += bytesRead;
306 if (offset >= sizeof(buffer)) {
307 err = ERROR_MORE_DATA;
308 }
309 } else {
310 err = GetLastError();
311 }
312 }
313
314 if (data.readPipe) {
315 CloseHandle(data.readPipe);
316 }
317
318 if (hThread) {
319 // Allow the thread some time to clean up
320 int thread_err;
321 switch (WaitForSingleObject(hThread, 100)) {
322 case WAIT_OBJECT_0:
323 // Thread ended cleanly
324 if (!GetExitCodeThread(hThread, (LPDWORD)&thread_err)) {
325 thread_err = GetLastError();
326 }
327 break;
328 case WAIT_TIMEOUT:
329 // Probably stuck - there's not much we can do, unfortunately
330 thread_err = WAIT_TIMEOUT;
331 break;
332 default:
333 thread_err = GetLastError();
334 break;
335 }
336 // An error on our side is more likely to be relevant than one from
337 // the thread, but if we don't have one on our side we'll take theirs.
338 if (err == 0 || err == ERROR_BROKEN_PIPE) {
339 err = thread_err;
340 }
341
342 CloseHandle(hThread);
343 }
344
345 CloseHandle(data.initEvent);
346 CloseHandle(data.connectEvent);
347 hThread = NULL;
348
349 Py_END_ALLOW_THREADS
350
351 PyMem_Free((void *)data.query);
352
353 if (err == ERROR_MORE_DATA) {
354 PyErr_Format(PyExc_OSError, "Query returns more than %zd characters", Py_ARRAY_LENGTH(buffer));
355 return NULL;
356 } else if (err) {
357 PyErr_SetFromWindowsErr(err);
358 return NULL;
359 }
360
361 if (!offset) {
362 return PyUnicode_FromStringAndSize(NULL, 0);
363 }
364 return PyUnicode_FromWideChar(buffer, offset / sizeof(buffer[0]) - 1);
365 }
366
367
368 static PyMethodDef wmi_functions[] = {
369 _WMI_EXEC_QUERY_METHODDEF
370 { NULL, NULL, 0, NULL }
371 };
372
373 static PyModuleDef_Slot wmi_slots[] = {
374 {Py_mod_gil, Py_MOD_GIL_NOT_USED},
375 {0, NULL},
376 };
377
378 static PyModuleDef wmi_def = {
379 PyModuleDef_HEAD_INIT,
380 "_wmi",
381 NULL, // doc
382 0, // m_size
383 wmi_functions, // m_methods
384 wmi_slots, // m_slots
385 };
386
387 extern "C" {
PyInit__wmi(void)388 PyMODINIT_FUNC PyInit__wmi(void)
389 {
390 return PyModuleDef_Init(&wmi_def);
391 }
392 }
393