winevent-tracker

cargo.toml

[package]
name = "winevent-tracker"
version = "0.1.0"
edition = "2021"

[dependencies]
clearscreen = "3"
windows = { version = "0.54", features = [
    "Win32_Foundation",
    "Win32_System_Threading",
    "Win32_UI_Accessibility",
    "Win32_UI_WindowsAndMessaging",
] }

main.rs

use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::sync::Mutex;
use std::sync::OnceLock;
use std::thread;
use std::time::Duration;
use windows::core::PWSTR;
use windows::Win32::Foundation::HANDLE;
use windows::Win32::Foundation::HWND;
use windows::Win32::System::Threading::OpenProcess;
use windows::Win32::System::Threading::PROCESS_NAME_WIN32;
use windows::Win32::System::Threading::PROCESS_QUERY_INFORMATION;
use windows::Win32::System::Threading::QueryFullProcessImageNameW;
use windows::Win32::UI::Accessibility::HWINEVENTHOOK;
use windows::Win32::UI::Accessibility::SetWinEventHook;
use windows::Win32::UI::WindowsAndMessaging::DispatchMessageW;
use windows::Win32::UI::WindowsAndMessaging::GetWindowThreadProcessId;
use windows::Win32::UI::WindowsAndMessaging::GetMessageW;
use windows::Win32::UI::WindowsAndMessaging::TranslateMessage;
use windows::Win32::UI::WindowsAndMessaging::EVENT_MAX;
use windows::Win32::UI::WindowsAndMessaging::EVENT_MIN;
use windows::Win32::UI::WindowsAndMessaging::MSG;


static STATE: OnceLock<Mutex<HashMap<String, isize>>> = OnceLock::new();
static HWND_MAP: OnceLock<Mutex<HashMap<isize, String>>> = OnceLock::new();

pub fn exe_path(handle: HANDLE) -> String {
    let mut len = 260_u32;
    let mut path: Vec<u16> = vec![0; len as usize];
    let text_ptr = path.as_mut_ptr();

    unsafe {
        if let Err(error) =
            QueryFullProcessImageNameW(handle, PROCESS_NAME_WIN32, PWSTR(text_ptr), &mut len) {
            println!("{error}");
        }
    }

    String::from_utf16(&path[..len as usize]).unwrap_or_default()
}

pub fn window_thread_process_id(hwnd: HWND) -> (u32, u32) {
    let mut process_id: u32 = 0;

    // Behaviour is undefined if an invalid HWND is given
    // https://docs.microsoft.com/en-us/windows/win32/api/winuser/nf-winuser-getwindowthreadprocessid
    let thread_id = unsafe {
        GetWindowThreadProcessId(hwnd, Option::from(std::ptr::addr_of_mut!(process_id)))
    };

    (process_id, thread_id)
}

pub fn process_handle(process_id: u32) -> Option<HANDLE> {
    unsafe { OpenProcess(PROCESS_QUERY_INFORMATION, false, process_id) }.ok()
}

pub extern "system" fn win_event_hook(
    _h_win_event_hook: HWINEVENTHOOK,
    event: u32,
    hwnd: HWND,
    id_object: i32,
    _id_child: i32,
    _id_event_thread: u32,
    _dwms_event_time: u32,
) {
    // OBJID_WINDOW
    if id_object != 0 {
        return;
    }

    let mut state = STATE
        .get_or_init(|| Mutex::new(HashMap::new()))
        .lock().unwrap();

    let mut hwnd_map = HWND_MAP
        .get_or_init(|| Mutex::new(HashMap::new()))
        .lock().unwrap();

    let exe_path = {
        if let Some(exe) = hwnd_map.get(&hwnd.0) {
            exe.clone()
        } else {
            let handle = process_handle(window_thread_process_id(hwnd).0);
            let mut exe = String::new();
            if let Some(handle) = handle {
                let exe_path = exe_path(handle);
                hwnd_map.insert(hwnd.0, exe_path.clone());
                exe = exe_path.clone();
            }

            if exe.is_empty() {
                return;
            } else {
                exe
            }
        }
    };

    match event {
        // create
        0x8000 => {
            *state.entry(exe_path).or_insert(0) += 1;
        }
        // destroy
        0x8001 => {
            let mut remove = false;
            if let Entry::Occupied(entry) = state.entry(exe_path.clone()) {
                if *entry.get() == 1 {
                    remove = true;
                } else {
                    *entry.into_mut() -= 1;
                }
            }

            if remove {
                state.remove(&exe_path);
            }
        }
        _ => {}
    }

    let _ = clearscreen::clear();

    dbg!(state);
}

fn main() {
    thread::spawn(move || {
        unsafe {
            SetWinEventHook(
                EVENT_MIN,
                EVENT_MAX,
                None,
                Some(win_event_hook),
                0,
                0,
                0,
            )
        };

        loop {
            let mut msg: MSG = MSG::default();
            unsafe {
                if !GetMessageW(&mut msg, HWND(0), 0, 0).as_bool() {
                    println!("windows event processing shutdown");
                    break;
                };
                TranslateMessage(&msg);
                DispatchMessageW(&msg);
            }
        }
    });

    loop {
        thread::sleep(Duration::from_secs(1));
    }
}