ohadravid / wmi-rs

WMI crate for rust
Apache License 2.0
84 stars 27 forks source link

Helper to resolve registry links #86

Open Enyium opened 11 months ago

Enyium commented 11 months ago

You already have the helper wmi::query::quote_and_escape_wql_str(). This is about adding another helper, specifically to transform a registry path like HKEY_CURRENT_USER\SOFTWARE\Foo\Bar into HKEY_USERS\<SID>\SOFTWARE\Foo\Bar. This is necessary to make the WMI query SELECT * FROM RegistryValueChangeEvent WHERE ... work. I'm sure it's also necessary for other registry-related queries.

I wrote an implementation for this - not as a helper function, but as an iterator, so that the SID has to be found out only once (in my case, when listening for changes to multiple registry values):

pub trait MakeResolveRegKeysIter<'a, I>
where
    I: IntoIterator<Item = (&'a str, &'a str)>,
{
    /// Creates an iterator that yields resolved hive-subkey pairs.
    ///
    /// This means changing `("HKEY_CURRENT_USER", r"foo\bar")` to `("HKEY_USERS", r"<SID>\foo\bar")` and passing through any other pairs unchanged. [`HKEY_CLASSES_ROOT`](https://learn.microsoft.com/en-us/windows/win32/sysinfo/hkey-classes-root-key) links to `HKEY_LOCAL_MACHINE\SOFTWARE\Classes` as well as `HKEY_CURRENT_USER\SOFTWARE\Classes` in a merging way, which is why it can't be resolved.
    ///
    /// The subkey shouldn't be preceded by a backslash.
    fn resolve_reg_keys(self) -> ResolveRegKeys<'a, I::IntoIter>;
}

impl<'a, I> MakeResolveRegKeysIter<'a, I::IntoIter> for I
where
    I: IntoIterator<Item = (&'a str, &'a str)>,
{
    fn resolve_reg_keys(self) -> ResolveRegKeys<'a, I::IntoIter> {
        ResolveRegKeys {
            iter: self.into_iter(),
            current_user_sid: None,
        }
    }
}

/// The iterator returned by [`MakeResolveRegKeysIter::resolve_reg_keys()`].
pub struct ResolveRegKeys<'a, I>
where
    I: Iterator<Item = (&'a str, &'a str)>,
{
    iter: I,
    current_user_sid: Option<String>,
}

impl<'a, I> Iterator for ResolveRegKeys<'a, I>
where
    I: Iterator<Item = (&'a str, &'a str)>,
{
    type Item = WMIResult<(&'a str, Cow<'a, str>)>;

    fn next(&mut self) -> Option<Self::Item> {
        Some(Ok(match self.iter.next()? {
            ("HKEY_CURRENT_USER", subkey) => {
                let current_user_sid = if let Some(sid) = self.current_user_sid.as_ref() {
                    sid
                } else {
                    let sid = match current_user_sid() {
                        Ok(sid) => sid,
                        Err(error) => {
                            return Some(Err(WMIError::HResultError {
                                hres: error.code().0,
                            }))
                        }
                    };
                    self.current_user_sid.insert(sid)
                };

                (
                    "HKEY_USERS",
                    Cow::Owned(current_user_sid.to_owned() + r"\" + subkey),
                )
            }
            (hkey, subkey) => (hkey, Cow::Borrowed(subkey)),
        }))
    }
}

fn current_user_sid() -> windows::core::Result<String> {
    let process_token_handle = ResGuard::with_mut_acq_and_close_handle(|handle| unsafe {
        OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, handle).ok()
    })?;

    let mut sid_and_attrs_buffer = Vec::<u8>::new();
    let mut sid_and_attrs_buffer_size = 0;

    dual_call(
        FirstCallExpectation::Win32Error(ERROR_INSUFFICIENT_BUFFER),
        |getting_buffer_size| unsafe {
            GetTokenInformation(
                *process_token_handle,
                TokenUser,
                (!getting_buffer_size).then(|| {
                    sid_and_attrs_buffer.resize(sid_and_attrs_buffer_size as _, 0);
                    sid_and_attrs_buffer.as_mut_ptr().cast()
                }),
                sid_and_attrs_buffer_size,
                &mut sid_and_attrs_buffer_size,
            )
            .ok()
        },
    )?;

    let string_sid = unsafe {
        ResGuard::with_mut_pwstr_acq_and_local_free(|pwstr| {
            ConvertSidToStringSidW(
                (&*sid_and_attrs_buffer.as_ptr().cast::<SID_AND_ATTRIBUTES>()).Sid,
                pwstr,
            )
            .ok()
        })?
        .to_string()?
    };

    Ok(string_sid)
}

(Note that, in this version, I made use of my crate windows-helpers.)

Are you willing to take such an iterator into your repo? If you want, I can provide it as a runnable, cloneable repo for testing or intermediate development purposes.