use crate::de::wbem_class_de::from_wbem_class_obj;
use crate::{
BStr,
connection::WMIConnection, safearray::safe_array_to_vec_of_strings, utils::check_hres, Variant,
WMIError,
};
use log::trace;
use serde::de;
use std::convert::TryInto;
use std::{mem, ptr, ptr::NonNull};
use winapi::um::oaidl::VARIANT;
use winapi::um::oleauto::VariantClear;
use winapi::{
shared::ntdef::NULL,
um::{
oaidl::SAFEARRAY,
oleauto::SafeArrayDestroy,
wbemcli::{
IEnumWbemClassObject, IWbemClassObject, WBEM_FLAG_ALWAYS, WBEM_FLAG_NONSYSTEM_ONLY,
WBEM_INFINITE,
},
},
};
#[derive(Debug)]
pub struct IWbemClassWrapper {
pub inner: NonNull<IWbemClassObject>,
}
impl IWbemClassWrapper {
pub unsafe fn new(ptr: NonNull<IWbemClassObject>) -> Self {
Self { inner: ptr }
}
pub unsafe fn clone(ptr: NonNull<IWbemClassObject>) -> Self {
let refcount = ptr.as_ref().AddRef();
trace!("Reference count: {}", refcount);
Self::new(ptr)
}
pub fn list_properties(&self) -> Result<Vec<String>, WMIError> {
let mut p_names = NULL as *mut SAFEARRAY;
let ptr = self.inner.as_ptr();
unsafe {
check_hres((*ptr).GetNames(
ptr::null(),
(WBEM_FLAG_ALWAYS | WBEM_FLAG_NONSYSTEM_ONLY) as i32,
ptr::null_mut(),
&mut p_names,
))?;
let res = safe_array_to_vec_of_strings(p_names);
check_hres(SafeArrayDestroy(p_names))?;
res
}
}
pub fn get_property(&self, property_name: &str) -> Result<Variant, WMIError> {
let name_prop = BStr::from_str(property_name)?;
let mut vt_prop: VARIANT = unsafe { mem::zeroed() };
unsafe {
check_hres((*self.inner.as_ptr()).Get(
name_prop.as_lpcwstr(),
0,
&mut vt_prop,
ptr::null_mut(),
ptr::null_mut(),
))?;
let property_value = Variant::from_variant(vt_prop)?;
check_hres(VariantClear(&mut vt_prop))?;
Ok(property_value)
}
}
pub fn path(&self) -> Result<String, WMIError> {
self.get_property("__Path").and_then(Variant::try_into)
}
pub fn class(&self) -> Result<String, WMIError> {
self.get_property("__Class").and_then(Variant::try_into)
}
pub fn into_desr<T>(self) -> Result<T, WMIError>
where
T: de::DeserializeOwned,
{
from_wbem_class_obj(&self).map_err(WMIError::from)
}
}
impl Drop for IWbemClassWrapper {
fn drop(&mut self) {
let ptr = self.inner.as_ptr();
unsafe {
(*ptr).Release();
}
}
}
pub struct QueryResultEnumerator<'a> {
_wmi_con: &'a WMIConnection,
p_enumerator: Option<NonNull<IEnumWbemClassObject>>,
}
impl<'a> QueryResultEnumerator<'a> {
pub unsafe fn new(wmi_con: &'a WMIConnection, p_enumerator: *mut IEnumWbemClassObject) -> Self {
Self {
_wmi_con: wmi_con,
p_enumerator: NonNull::new(p_enumerator),
}
}
}
impl<'a> Drop for QueryResultEnumerator<'a> {
fn drop(&mut self) {
if let Some(p_enumerator) = self.p_enumerator {
unsafe {
(*p_enumerator.as_ptr()).Release();
}
}
}
}
impl<'a> Iterator for QueryResultEnumerator<'a> {
type Item = Result<IWbemClassWrapper, WMIError>;
fn next(&mut self) -> Option<Self::Item> {
let mut pcls_obj = NULL as *mut IWbemClassObject;
let mut return_value = 0;
let raw_enumerator_prt = self.p_enumerator?.as_ptr();
let res = unsafe {
check_hres((*raw_enumerator_prt).Next(
WBEM_INFINITE as i32,
1,
&mut pcls_obj,
&mut return_value,
))
};
if let Err(e) = res {
return Some(Err(e));
}
if return_value == 0 {
return None;
}
trace!(
"Got enumerator {:?} and obj {:?}",
self.p_enumerator,
pcls_obj
);
let pcls_ptr = NonNull::new(pcls_obj).ok_or(WMIError::NullPointerResult);
match pcls_ptr {
Err(e) => Some(Err(e)),
Ok(pcls_ptr) => Some(Ok( unsafe { IWbemClassWrapper::new(pcls_ptr) })),
}
}
}