ohadravid / wmi-rs

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

Is there a way to suscribe to multiple events in the same thread? #78

Closed sn99 closed 1 year ago

sn99 commented 1 year ago

I am tying to subscribe to process and thread event at the same time, ideally the code would look something like this (with commented parts in):

use futures::StreamExt;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Duration;

use wmi::FilterValue;
use wmi::{COMLibrary, WMIConnection};

#[derive(Deserialize, Debug)]
struct __InstanceCreationEvent {
    TargetInstance: Win32_Process,
}

#[derive(Deserialize, Debug)]
struct Win32_Process {
    ProcessID: u32,
}

#[derive(Deserialize, Serialize, Debug)]
struct __InstanceCreationThreadEvent {
    TargetInstance: Win32_Thread,
}

#[derive(Deserialize, Serialize, Debug)]
struct Win32_Thread {
    ProcessHandle: String,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let com_con = COMLibrary::new()?;
    let wmi = WMIConnection::new(com_con)?;

    let mut filters_process = HashMap::new();
    let mut filters_thread = HashMap::new();

    filters_process.insert(
        "TargetInstance".to_owned(),
        FilterValue::is_a::<Win32_Process>()?,
    );
    filters_thread.insert(
        "TargetInstance".to_owned(),
        FilterValue::is_a::<Win32_Thread>()?,
    );

    let mut stream_process = wmi.async_filtered_notification::<__InstanceCreationEvent>(
        &filters_process,
        Some(Duration::from_secs(1)),
    )?;

    /*
    let wmi_thread = WMIConnection::new(com_con)?;

    let mut stream_thread = wmi.async_filtered_notification::<__InstanceCreationThreadEvent>(
        &filters_thread,
        Some(Duration::from_secs(1)),
    )?;
    */

    while let Some(k) = stream_process.next().await {
        println!("{:?}", k);
    }

    Ok(())
}
ohadravid commented 1 year ago

Hi @sn99, I think you want to use #serde(rename = "...") with tokio::select!.

Here's an example which does what I think you want 🔍:

use futures::StreamExt;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Duration;

use wmi::FilterValue;
use wmi::{COMLibrary, WMIConnection};

#[derive(Deserialize, Debug)]
struct __InstanceCreationEvent {
    TargetInstance: Win32_Process,
}

#[derive(Deserialize, Debug)]
struct Win32_Process {
    ProcessID: u32,
}

#[derive(Deserialize, Serialize, Debug)]
#[serde(rename = "__InstanceCreationEvent")]
struct __InstanceCreationThreadEvent {
    TargetInstance: Win32_Thread,
}

#[derive(Deserialize, Serialize, Debug)]
struct Win32_Thread {
    ProcessHandle: String,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let com_con = COMLibrary::new()?;
    let wmi = WMIConnection::new(com_con)?;

    let mut filters_process = HashMap::new();
    let mut filters_thread = HashMap::new();

    filters_process.insert(
        "TargetInstance".to_owned(),
        FilterValue::is_a::<Win32_Process>()?,
    );
    filters_thread.insert(
        "TargetInstance".to_owned(),
        FilterValue::is_a::<Win32_Thread>()?,
    );

    let mut stream_process = wmi.async_filtered_notification::<__InstanceCreationEvent>(
        &filters_process,
        Some(Duration::from_secs(1)),
    )?;
    let mut stream_thread = wmi.async_filtered_notification::<__InstanceCreationThreadEvent>(
        &filters_thread,
        Some(Duration::from_secs(1)),
    )?;

    loop {
        tokio::select! {
          Some(proc) = stream_process.next() => { dbg!(proc); },
          Some(thread) = stream_thread.next() => { dbg!(thread); },
          else => break,
        }
    }
    Ok(())
}
ohadravid commented 1 year ago

We also (partially) support enum serialization, but it's missing for fields. I'll add this later, but it will be possible to do:

#[derive(Deserialize, Debug)]
struct Win32_Process {
    ProcessID: u32,
}

#[derive(Deserialize, Debug)]
struct Win32_Thread {
    ProcessHandle: String,
}

#[derive(Deserialize, Debug)]
enum Instance {
    #[serde(rename = "Win32_Process")]
    Process(Win32_Process),
    #[serde(rename = "Win32_Thread")]
    Thread(Win32_Thread),
}

#[derive(Deserialize, Debug)]
struct __InstanceCreationEvent {
    TargetInstance: Instance,
}

Update: see https://github.com/ohadravid/wmi-rs/pull/79

sn99 commented 1 year ago

Yep @ohadravid this is exactly what I was looking for #79

ohadravid commented 1 year ago

Closed as #79 is now merged