mirror of
https://github.com/python/cpython.git
synced 2025-08-04 00:48:58 +00:00
gh-89545: Updates platform module to use new internal _wmi module on Windows to directly query OS properties (GH-96289)
This commit is contained in:
Steve Dower
GitHub
parent
4114bcc9ef
commit
de33df27aa
17 changed files with 856 additions and 88 deletions
307
PC/_wmimodule.cpp
Normal file
307
PC/_wmimodule.cpp
Normal file
|
|
@ -0,0 +1,307 @@
|
|||
//
|
||||
// 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_APARTMENTTHREADED);
|
||||
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 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;
|
||||
}
|
||||
|
||||
|
||||
/*[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[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
|
||||
|
||||
if (!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);
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
// Allow the thread some time to clean up
|
||||
switch (WaitForSingleObject(hThread, 1000)) {
|
||||
case WAIT_OBJECT_0:
|
||||
// Thread ended cleanly
|
||||
if (!GetExitCodeThread(hThread, (LPDWORD)&err)) {
|
||||
err = GetLastError();
|
||||
}
|
||||
break;
|
||||
case WAIT_TIMEOUT:
|
||||
// Probably stuck - there's not much we can do, unfortunately
|
||||
if (err == 0 || err == ERROR_BROKEN_PIPE) {
|
||||
err = WAIT_TIMEOUT;
|
||||
}
|
||||
break;
|
||||
default:
|
||||
if (err == 0 || err == ERROR_BROKEN_PIPE) {
|
||||
err = GetLastError();
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
CloseHandle(hThread);
|
||||
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 wmi_def = {
|
||||
PyModuleDef_HEAD_INIT,
|
||||
"_wmi",
|
||||
NULL, // doc
|
||||
0, // m_size
|
||||
wmi_functions
|
||||
};
|
||||
|
||||
extern "C" {
|
||||
PyMODINIT_FUNC PyInit__wmi(void)
|
||||
{
|
||||
return PyModuleDef_Init(&wmi_def);
|
||||
}
|
||||
}
|
||||
75
PC/clinic/_wmimodule.cpp.h
Normal file
75
PC/clinic/_wmimodule.cpp.h
Normal file
|
|
@ -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]*/
|
||||
Loading…
Add table
Add a link
Reference in a new issue