claudetools/projects/msp-tools/guru-connect/agent/src/viewer/input.rs

//! Low-level keyboard hook for capturing all keys including Win key

use super::InputEvent;
#[cfg(windows)]
use crate::proto;
use anyhow::Result;
use tokio::sync::mpsc;
#[cfg(windows)]
use tracing::trace;

#[cfg(windows)]
use windows::{
    Win32::Foundation::{LPARAM, LRESULT, WPARAM},
    Win32::UI::WindowsAndMessaging::{
        CallNextHookEx, DispatchMessageW, GetMessageW, PeekMessageW, SetWindowsHookExW,
        TranslateMessage, UnhookWindowsHookEx, HHOOK, KBDLLHOOKSTRUCT, MSG, PM_REMOVE,
        WH_KEYBOARD_LL, WM_KEYDOWN, WM_KEYUP, WM_SYSKEYDOWN, WM_SYSKEYUP,
    },
};

#[cfg(windows)]
use std::sync::OnceLock;

#[cfg(windows)]
static INPUT_TX: OnceLock<mpsc::Sender<InputEvent>> = OnceLock::new();

#[cfg(windows)]
static mut HOOK_HANDLE: HHOOK = HHOOK(std::ptr::null_mut());

/// Virtual key codes for special keys
#[cfg(windows)]
mod vk {
    pub const VK_LWIN: u32 = 0x5B;
    pub const VK_RWIN: u32 = 0x5C;
    pub const VK_APPS: u32 = 0x5D;
    pub const VK_LSHIFT: u32 = 0xA0;
    pub const VK_RSHIFT: u32 = 0xA1;
    pub const VK_LCONTROL: u32 = 0xA2;
    pub const VK_RCONTROL: u32 = 0xA3;
    pub const VK_LMENU: u32 = 0xA4;  // Left Alt
    pub const VK_RMENU: u32 = 0xA5;  // Right Alt
    pub const VK_TAB: u32 = 0x09;
    pub const VK_ESCAPE: u32 = 0x1B;
    pub const VK_SNAPSHOT: u32 = 0x2C;  // Print Screen
}

#[cfg(windows)]
pub struct KeyboardHook {
    _hook: HHOOK,
}

#[cfg(windows)]
impl KeyboardHook {
    pub fn new(input_tx: mpsc::Sender<InputEvent>) -> Result<Self> {
        // Store the sender globally for the hook callback
        INPUT_TX.set(input_tx).map_err(|_| anyhow::anyhow!("Input TX already set"))?;

        unsafe {
            let hook = SetWindowsHookExW(
                WH_KEYBOARD_LL,
                Some(keyboard_hook_proc),
                None,
                0,
            )?;

            HOOK_HANDLE = hook;
            Ok(Self { _hook: hook })
        }
    }
}

#[cfg(windows)]
impl Drop for KeyboardHook {
    fn drop(&mut self) {
        unsafe {
            if !HOOK_HANDLE.0.is_null() {
                let _ = UnhookWindowsHookEx(HOOK_HANDLE);
                HOOK_HANDLE = HHOOK(std::ptr::null_mut());
            }
        }
    }
}

#[cfg(windows)]
unsafe extern "system" fn keyboard_hook_proc(
    code: i32,
    wparam: WPARAM,
    lparam: LPARAM,
) -> LRESULT {
    if code >= 0 {
        let kb_struct = &*(lparam.0 as *const KBDLLHOOKSTRUCT);
        let vk_code = kb_struct.vkCode;
        let scan_code = kb_struct.scanCode;

        let is_down = wparam.0 as u32 == WM_KEYDOWN || wparam.0 as u32 == WM_SYSKEYDOWN;
        let is_up = wparam.0 as u32 == WM_KEYUP || wparam.0 as u32 == WM_SYSKEYUP;

        if is_down || is_up {
            // Check if this is a key we want to intercept (Win key, Alt+Tab, etc.)
            let should_intercept = matches!(
                vk_code,
                vk::VK_LWIN | vk::VK_RWIN | vk::VK_APPS
            );

            // Send the key event to the remote
            if let Some(tx) = INPUT_TX.get() {
                let event = proto::KeyEvent {
                    down: is_down,
                    key_type: proto::KeyEventType::KeyVk as i32,
                    vk_code,
                    scan_code,
                    unicode: String::new(),
                    modifiers: Some(get_current_modifiers()),
                };

                let _ = tx.try_send(InputEvent::Key(event));
                trace!("Key hook: vk={:#x} scan={} down={}", vk_code, scan_code, is_down);
            }

            // For Win key, consume the event so it doesn't open Start menu locally
            if should_intercept {
                return LRESULT(1);
            }
        }
    }

    CallNextHookEx(HOOK_HANDLE, code, wparam, lparam)
}

#[cfg(windows)]
fn get_current_modifiers() -> proto::Modifiers {
    use windows::Win32::UI::Input::KeyboardAndMouse::GetAsyncKeyState;

    unsafe {
        proto::Modifiers {
            ctrl: GetAsyncKeyState(0x11) < 0,   // VK_CONTROL
            alt: GetAsyncKeyState(0x12) < 0,    // VK_MENU
            shift: GetAsyncKeyState(0x10) < 0,  // VK_SHIFT
            meta: GetAsyncKeyState(0x5B) < 0 || GetAsyncKeyState(0x5C) < 0, // VK_LWIN/RWIN
            caps_lock: GetAsyncKeyState(0x14) & 1 != 0, // VK_CAPITAL
            num_lock: GetAsyncKeyState(0x90) & 1 != 0,  // VK_NUMLOCK
        }
    }
}

/// Pump Windows message queue (required for hooks to work)
#[cfg(windows)]
pub fn pump_messages() {
    unsafe {
        let mut msg = MSG::default();
        while PeekMessageW(&mut msg, None, 0, 0, PM_REMOVE).as_bool() {
            let _ = TranslateMessage(&msg);
            DispatchMessageW(&msg);
        }
    }
}

// Non-Windows stubs
#[cfg(not(windows))]
#[allow(dead_code)]
pub struct KeyboardHook;

#[cfg(not(windows))]
#[allow(dead_code)]
impl KeyboardHook {
    pub fn new(_input_tx: mpsc::Sender<InputEvent>) -> Result<Self> {
        Ok(Self)
    }
}

#[cfg(not(windows))]
#[allow(dead_code)]
pub fn pump_messages() {}