From 720fd7f63c81392d888facc31f3645e6197b9ddd Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sun, 15 Feb 2026 23:41:23 +0900 Subject: [PATCH] Add _wmi module Implement _wmi.exec_query via COM/WMI using pipe+thread pattern matching _wmimodule.cpp. Uses raw COM vtable calls for IWbemLocator, IWbemServices, IEnumWbemClassObject, and IWbemClassObject interfaces. --- crates/vm/src/stdlib/_wmi.rs | 698 +++++++++++++++++++++++++++++++++++ crates/vm/src/stdlib/mod.rs | 4 + 2 files changed, 702 insertions(+) create mode 100644 crates/vm/src/stdlib/_wmi.rs diff --git a/crates/vm/src/stdlib/_wmi.rs b/crates/vm/src/stdlib/_wmi.rs new file mode 100644 index 000000000..f25f4b23b --- /dev/null +++ b/crates/vm/src/stdlib/_wmi.rs @@ -0,0 +1,698 @@ +// spell-checker:disable +#![allow(non_snake_case)] + +pub(crate) use _wmi::module_def; + +// COM/WMI FFI declarations (not inside pymodule to avoid macro issues) +mod wmi_ffi { + #![allow(unsafe_op_in_unsafe_fn)] + use core::ffi::c_void; + + pub type HRESULT = i32; + + #[repr(C)] + pub struct GUID { + pub data1: u32, + pub data2: u16, + pub data3: u16, + pub data4: [u8; 8], + } + + // Opaque VARIANT type (24 bytes covers both 32-bit and 64-bit) + #[repr(C, align(8))] + pub struct VARIANT([u64; 3]); + + impl VARIANT { + pub fn zeroed() -> Self { + VARIANT([0u64; 3]) + } + } + + // CLSID_WbemLocator = {4590F811-1D3A-11D0-891F-00AA004B2E24} + pub const CLSID_WBEM_LOCATOR: GUID = GUID { + data1: 0x4590F811, + data2: 0x1D3A, + data3: 0x11D0, + data4: [0x89, 0x1F, 0x00, 0xAA, 0x00, 0x4B, 0x2E, 0x24], + }; + + // IID_IWbemLocator = {DC12A687-737F-11CF-884D-00AA004B2E24} + pub const IID_IWBEM_LOCATOR: GUID = GUID { + data1: 0xDC12A687, + data2: 0x737F, + data3: 0x11CF, + data4: [0x88, 0x4D, 0x00, 0xAA, 0x00, 0x4B, 0x2E, 0x24], + }; + + // COM constants + pub const COINIT_APARTMENTTHREADED: u32 = 0x2; + pub const CLSCTX_INPROC_SERVER: u32 = 0x1; + pub const RPC_C_AUTHN_LEVEL_DEFAULT: u32 = 0; + pub const RPC_C_IMP_LEVEL_IMPERSONATE: u32 = 3; + pub const RPC_C_AUTHN_LEVEL_CALL: u32 = 3; + pub const RPC_C_AUTHN_WINNT: u32 = 10; + pub const RPC_C_AUTHZ_NONE: u32 = 0; + pub const EOAC_NONE: u32 = 0; + pub const RPC_E_TOO_LATE: HRESULT = 0x80010119_u32 as i32; + + // WMI constants + pub const WBEM_FLAG_FORWARD_ONLY: i32 = 0x20; + pub const WBEM_FLAG_RETURN_IMMEDIATELY: i32 = 0x10; + pub const WBEM_S_FALSE: HRESULT = 1; + pub const WBEM_S_NO_MORE_DATA: HRESULT = 0x40005; + pub const WBEM_INFINITE: i32 = -1; + pub const WBEM_FLAVOR_MASK_ORIGIN: i32 = 0x60; + pub const WBEM_FLAVOR_ORIGIN_SYSTEM: i32 = 0x40; + + #[link(name = "ole32")] + unsafe extern "system" { + pub fn CoInitializeEx(pvReserved: *mut c_void, dwCoInit: u32) -> HRESULT; + pub fn CoUninitialize(); + pub fn CoInitializeSecurity( + pSecDesc: *const c_void, + cAuthSvc: i32, + asAuthSvc: *const c_void, + pReserved1: *const c_void, + dwAuthnLevel: u32, + dwImpLevel: u32, + pAuthList: *const c_void, + dwCapabilities: u32, + pReserved3: *const c_void, + ) -> HRESULT; + pub fn CoCreateInstance( + rclsid: *const GUID, + pUnkOuter: *mut c_void, + dwClsContext: u32, + riid: *const GUID, + ppv: *mut *mut c_void, + ) -> HRESULT; + pub fn CoSetProxyBlanket( + pProxy: *mut c_void, + dwAuthnSvc: u32, + dwAuthzSvc: u32, + pServerPrincName: *const u16, + dwAuthnLevel: u32, + dwImpLevel: u32, + pAuthInfo: *const c_void, + dwCapabilities: u32, + ) -> HRESULT; + } + + #[link(name = "oleaut32")] + unsafe extern "system" { + pub fn SysAllocString(psz: *const u16) -> *mut u16; + pub fn SysFreeString(bstrString: *mut u16); + pub fn VariantClear(pvarg: *mut VARIANT) -> HRESULT; + } + + #[link(name = "propsys")] + unsafe extern "system" { + pub fn VariantToString(varIn: *const VARIANT, pszBuf: *mut u16, cchBuf: u32) -> HRESULT; + } + + /// Release a COM object (IUnknown::Release, vtable index 2) + pub unsafe fn com_release(this: *mut c_void) { + if !this.is_null() { + let vtable = *(this as *const *const usize); + let release: unsafe extern "system" fn(*mut c_void) -> u32 = + core::mem::transmute(*vtable.add(2)); + release(this); + } + } + + /// IWbemLocator::ConnectServer (vtable index 3) + #[allow(clippy::too_many_arguments)] + pub unsafe fn locator_connect_server( + this: *mut c_void, + network_resource: *const u16, + user: *const u16, + password: *const u16, + locale: *const u16, + security_flags: i32, + authority: *const u16, + ctx: *mut c_void, + services: *mut *mut c_void, + ) -> HRESULT { + let vtable = *(this as *const *const usize); + let method: unsafe extern "system" fn( + *mut c_void, + *const u16, + *const u16, + *const u16, + *const u16, + i32, + *const u16, + *mut c_void, + *mut *mut c_void, + ) -> HRESULT = core::mem::transmute(*vtable.add(3)); + method( + this, + network_resource, + user, + password, + locale, + security_flags, + authority, + ctx, + services, + ) + } + + /// IWbemServices::ExecQuery (vtable index 20) + pub unsafe fn services_exec_query( + this: *mut c_void, + query_language: *const u16, + query: *const u16, + flags: i32, + ctx: *mut c_void, + enumerator: *mut *mut c_void, + ) -> HRESULT { + let vtable = *(this as *const *const usize); + let method: unsafe extern "system" fn( + *mut c_void, + *const u16, + *const u16, + i32, + *mut c_void, + *mut *mut c_void, + ) -> HRESULT = core::mem::transmute(*vtable.add(20)); + method(this, query_language, query, flags, ctx, enumerator) + } + + /// IEnumWbemClassObject::Next (vtable index 4) + pub unsafe fn enum_next( + this: *mut c_void, + timeout: i32, + count: u32, + objects: *mut *mut c_void, + returned: *mut u32, + ) -> HRESULT { + let vtable = *(this as *const *const usize); + let method: unsafe extern "system" fn( + *mut c_void, + i32, + u32, + *mut *mut c_void, + *mut u32, + ) -> HRESULT = core::mem::transmute(*vtable.add(4)); + method(this, timeout, count, objects, returned) + } + + /// IWbemClassObject::BeginEnumeration (vtable index 8) + pub unsafe fn object_begin_enumeration(this: *mut c_void, enum_flags: i32) -> HRESULT { + let vtable = *(this as *const *const usize); + let method: unsafe extern "system" fn(*mut c_void, i32) -> HRESULT = + core::mem::transmute(*vtable.add(8)); + method(this, enum_flags) + } + + /// IWbemClassObject::Next (vtable index 9) + pub unsafe fn object_next( + this: *mut c_void, + flags: i32, + name: *mut *mut u16, + val: *mut VARIANT, + cim_type: *mut i32, + flavor: *mut i32, + ) -> HRESULT { + let vtable = *(this as *const *const usize); + let method: unsafe extern "system" fn( + *mut c_void, + i32, + *mut *mut u16, + *mut VARIANT, + *mut i32, + *mut i32, + ) -> HRESULT = core::mem::transmute(*vtable.add(9)); + method(this, flags, name, val, cim_type, flavor) + } + + /// IWbemClassObject::EndEnumeration (vtable index 10) + pub unsafe fn object_end_enumeration(this: *mut c_void) -> HRESULT { + let vtable = *(this as *const *const usize); + let method: unsafe extern "system" fn(*mut c_void) -> HRESULT = + core::mem::transmute(*vtable.add(10)); + method(this) + } +} + +#[pymodule] +mod _wmi { + use super::wmi_ffi::*; + use crate::builtins::PyStrRef; + use crate::convert::ToPyException; + use crate::{PyResult, VirtualMachine}; + use core::ffi::c_void; + use core::ptr::{null, null_mut}; + use windows_sys::Win32::Foundation::{ + CloseHandle, ERROR_BROKEN_PIPE, ERROR_MORE_DATA, ERROR_NOT_ENOUGH_MEMORY, GetLastError, + HANDLE, WAIT_OBJECT_0, WAIT_TIMEOUT, + }; + use windows_sys::Win32::Storage::FileSystem::{ReadFile, WriteFile}; + use windows_sys::Win32::System::Pipes::CreatePipe; + use windows_sys::Win32::System::Threading::{ + CreateEventW, CreateThread, GetExitCodeThread, SetEvent, WaitForSingleObject, + }; + + const BUFFER_SIZE: usize = 8192; + + fn hresult_from_win32(err: u32) -> HRESULT { + if err == 0 { + 0 + } else { + ((err & 0xFFFF) | 0x80070000) as HRESULT + } + } + + fn succeeded(hr: HRESULT) -> bool { + hr >= 0 + } + + fn failed(hr: HRESULT) -> bool { + hr < 0 + } + + fn wide_str(s: &str) -> Vec { + s.encode_utf16().chain(core::iter::once(0)).collect() + } + + unsafe fn wcslen(s: *const u16) -> usize { + unsafe { + let mut len = 0; + while *s.add(len) != 0 { + len += 1; + } + len + } + } + + unsafe fn wait_event(event: HANDLE, timeout: u32) -> u32 { + unsafe { + match WaitForSingleObject(event, timeout) { + WAIT_OBJECT_0 => 0, + WAIT_TIMEOUT => WAIT_TIMEOUT, + _ => GetLastError(), + } + } + } + + struct QueryThreadData { + query: Vec, + write_pipe: HANDLE, + init_event: HANDLE, + connect_event: HANDLE, + } + + // SAFETY: QueryThreadData contains HANDLEs (isize) which are safe to send across threads + unsafe impl Send for QueryThreadData {} + + unsafe extern "system" fn query_thread(param: *mut c_void) -> u32 { + unsafe { query_thread_impl(param) } + } + + unsafe fn query_thread_impl(param: *mut c_void) -> u32 { + unsafe { + let data = Box::from_raw(param as *mut QueryThreadData); + let write_pipe = data.write_pipe; + let init_event = data.init_event; + let connect_event = data.connect_event; + + let mut locator: *mut c_void = null_mut(); + let mut services: *mut c_void = null_mut(); + let mut enumerator: *mut c_void = null_mut(); + let mut hr: HRESULT = 0; + + // gh-125315: Copy the query string first + let bstr_query = SysAllocString(data.query.as_ptr()); + if bstr_query.is_null() { + hr = hresult_from_win32(ERROR_NOT_ENOUGH_MEMORY); + } + + drop(data); + + if succeeded(hr) { + hr = CoInitializeEx(null_mut(), COINIT_APARTMENTTHREADED); + } + + if failed(hr) { + CloseHandle(write_pipe); + if !bstr_query.is_null() { + SysFreeString(bstr_query); + } + return hr as u32; + } + + 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. + if hr == RPC_E_TOO_LATE { + hr = 0; + } + + if succeeded(hr) { + hr = CoCreateInstance( + &CLSID_WBEM_LOCATOR, + null_mut(), + CLSCTX_INPROC_SERVER, + &IID_IWBEM_LOCATOR, + &mut locator, + ); + } + if succeeded(hr) && SetEvent(init_event) == 0 { + hr = hresult_from_win32(GetLastError()); + } + + if succeeded(hr) { + let root_cimv2 = wide_str("ROOT\\CIMV2"); + let bstr_root = SysAllocString(root_cimv2.as_ptr()); + hr = locator_connect_server( + locator, + bstr_root, + null(), + null(), + null(), + 0, + null(), + null_mut(), + &mut services, + ); + if !bstr_root.is_null() { + SysFreeString(bstr_root); + } + } + if succeeded(hr) && SetEvent(connect_event) == 0 { + 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) { + let wql = wide_str("WQL"); + let bstr_wql = SysAllocString(wql.as_ptr()); + hr = services_exec_query( + services, + bstr_wql, + bstr_query, + WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, + null_mut(), + &mut enumerator, + ); + if !bstr_wql.is_null() { + SysFreeString(bstr_wql); + } + } + + // Enumerate results and write to pipe + let mut value: *mut c_void; + let mut start_of_enum = true; + let null_sep: u16 = 0; + let eq_sign: u16 = b'=' as u16; + + while succeeded(hr) { + let mut got: u32 = 0; + let mut written: u32 = 0; + value = null_mut(); + hr = enum_next(enumerator, WBEM_INFINITE, 1, &mut value, &mut got); + + if hr == WBEM_S_FALSE { + hr = 0; + break; + } + if failed(hr) || got != 1 || value.is_null() { + continue; + } + + if !start_of_enum + && WriteFile( + write_pipe, + &null_sep as *const u16 as *const _, + 2, + &mut written, + null_mut(), + ) == 0 + { + hr = hresult_from_win32(GetLastError()); + com_release(value); + break; + } + start_of_enum = false; + + hr = object_begin_enumeration(value, 0); + if failed(hr) { + com_release(value); + break; + } + + while succeeded(hr) { + let mut prop_name: *mut u16 = null_mut(); + let mut prop_value = VARIANT::zeroed(); + let mut flavor: i32 = 0; + + hr = object_next( + value, + 0, + &mut prop_name, + &mut prop_value, + null_mut(), + &mut flavor, + ); + + if hr == WBEM_S_NO_MORE_DATA { + hr = 0; + break; + } + + if succeeded(hr) + && (flavor & WBEM_FLAVOR_MASK_ORIGIN) != WBEM_FLAVOR_ORIGIN_SYSTEM + { + let mut prop_str = [0u16; BUFFER_SIZE]; + hr = + VariantToString(&prop_value, prop_str.as_mut_ptr(), BUFFER_SIZE as u32); + + if succeeded(hr) { + let cb_str1 = (wcslen(prop_name) * 2) as u32; + let cb_str2 = (wcslen(prop_str.as_ptr()) * 2) as u32; + + if WriteFile( + write_pipe, + prop_name as *const _, + cb_str1, + &mut written, + null_mut(), + ) == 0 + || WriteFile( + write_pipe, + &eq_sign as *const u16 as *const _, + 2, + &mut written, + null_mut(), + ) == 0 + || WriteFile( + write_pipe, + prop_str.as_ptr() as *const _, + cb_str2, + &mut written, + null_mut(), + ) == 0 + || WriteFile( + write_pipe, + &null_sep as *const u16 as *const _, + 2, + &mut written, + null_mut(), + ) == 0 + { + hr = hresult_from_win32(GetLastError()); + } + } + + VariantClear(&mut prop_value); + SysFreeString(prop_name); + } + } + + object_end_enumeration(value); + com_release(value); + } + + // Cleanup + if !bstr_query.is_null() { + SysFreeString(bstr_query); + } + if !enumerator.is_null() { + com_release(enumerator); + } + if !services.is_null() { + com_release(services); + } + if !locator.is_null() { + com_release(locator); + } + CoUninitialize(); + CloseHandle(write_pipe); + + hr as u32 + } + } + + /// 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. + #[pyfunction] + fn exec_query(query: PyStrRef, vm: &VirtualMachine) -> PyResult { + let query_str = query.as_str(); + + if !query_str + .get(..7) + .is_some_and(|s| s.eq_ignore_ascii_case("select ")) + { + return Err(vm.new_value_error("only SELECT queries are supported".to_owned())); + } + + let query_wide = wide_str(query_str); + + let mut h_thread: HANDLE = null_mut(); + let mut err: u32 = 0; + let mut buffer = [0u16; BUFFER_SIZE]; + let mut offset: u32 = 0; + let mut bytes_read: u32 = 0; + + let mut read_pipe: HANDLE = null_mut(); + let mut write_pipe: HANDLE = null_mut(); + + unsafe { + let init_event = CreateEventW(null(), 1, 0, null()); + let connect_event = CreateEventW(null(), 1, 0, null()); + + if init_event.is_null() + || connect_event.is_null() + || CreatePipe(&mut read_pipe, &mut write_pipe, null(), 0) == 0 + { + err = GetLastError(); + } else { + let thread_data = Box::new(QueryThreadData { + query: query_wide, + write_pipe, + init_event, + connect_event, + }); + let thread_data_ptr = Box::into_raw(thread_data); + + h_thread = CreateThread( + null(), + 0, + Some(query_thread), + thread_data_ptr as *const _ as *mut _, + 0, + null_mut(), + ); + + if h_thread.is_null() { + err = GetLastError(); + // Thread didn't start, so recover data and close write pipe + let data = Box::from_raw(thread_data_ptr); + CloseHandle(data.write_pipe); + } + } + + // gh-112278: Timeout for COM init and WMI connection + if err == 0 { + err = wait_event(init_event, 1000); + if err == 0 { + err = wait_event(connect_event, 100); + } + } + + // Read results from pipe + while err == 0 { + let buf_ptr = (buffer.as_mut_ptr() as *mut u8).add(offset as usize); + let buf_remaining = (BUFFER_SIZE * 2) as u32 - offset; + + if ReadFile( + read_pipe, + buf_ptr as *mut _, + buf_remaining, + &mut bytes_read, + null_mut(), + ) != 0 + { + offset += bytes_read; + if offset >= (BUFFER_SIZE * 2) as u32 { + err = ERROR_MORE_DATA; + } + } else { + err = GetLastError(); + } + } + + if !read_pipe.is_null() { + CloseHandle(read_pipe); + } + + if !h_thread.is_null() { + let thread_err: u32; + match WaitForSingleObject(h_thread, 100) { + WAIT_OBJECT_0 => { + let mut exit_code: u32 = 0; + if GetExitCodeThread(h_thread, &mut exit_code) == 0 { + thread_err = GetLastError(); + } else { + thread_err = exit_code; + } + } + WAIT_TIMEOUT => { + thread_err = WAIT_TIMEOUT; + } + _ => { + thread_err = GetLastError(); + } + } + if err == 0 || err == ERROR_BROKEN_PIPE { + err = thread_err; + } + + CloseHandle(h_thread); + } + + CloseHandle(init_event); + CloseHandle(connect_event); + } + + if err == ERROR_MORE_DATA { + return Err(vm.new_os_error(format!( + "Query returns more than {} characters", + BUFFER_SIZE, + ))); + } else if err != 0 { + return Err(std::io::Error::from_raw_os_error(err as i32).to_pyexception(vm)); + } + + if offset == 0 { + return Ok(String::new()); + } + + let char_count = (offset as usize) / 2 - 1; + Ok(String::from_utf16_lossy(&buffer[..char_count])) + } +} diff --git a/crates/vm/src/stdlib/mod.rs b/crates/vm/src/stdlib/mod.rs index c9cd66346..5bec43f72 100644 --- a/crates/vm/src/stdlib/mod.rs +++ b/crates/vm/src/stdlib/mod.rs @@ -58,6 +58,8 @@ pub(crate) mod msvcrt; ))] mod pwd; +#[cfg(all(feature = "host_env", windows))] +mod _wmi; #[cfg(feature = "host_env")] pub(crate) mod signal; pub mod sys; @@ -129,5 +131,7 @@ pub fn builtin_module_defs(ctx: &Context) -> Vec<&'static PyModuleDef> { winapi::module_def(ctx), #[cfg(all(feature = "host_env", windows))] winreg::module_def(ctx), + #[cfg(all(feature = "host_env", windows))] + _wmi::module_def(ctx), ] }