moon-driver-utils 0.1.0

Windows Kernel Utils
use core::{
    ffi::c_void,
    ops::{Deref, DerefMut},
};

use alloc::string::String;
use wdk::println;
use wdk_sys::{
    ntddk::{
        ObReferenceObjectByHandle, ObfDereferenceObject, ObfReferenceObject, PsGetProcessId,
        PsLookupProcessByProcessId,
    },
    PsProcessType, NT_SUCCESS, PHANDLE, THREAD_ALL_ACCESS,
    _MODE::KernelMode,
};

use crate::{kernel_fucntion::PsGetProcessImageFileName, string::cstr_to_rust_str};

pub struct ProcessError {}

pub struct Process {
    raw: *mut c_void,
    pid: *mut c_void,
}

impl Process {
    pub fn from_raw(raw: *mut c_void) -> Result<Self, ProcessError> {
        let pid = unsafe { PsGetProcessId(raw as _) };
        unsafe {
            ObfReferenceObject(raw);
        }
        let r = Self { raw, pid };
        Ok(r)
    }

    pub fn from_pid(pid: *mut c_void) -> Result<Self, ProcessError> {
        let mut p = core::ptr::null_mut();
        unsafe {
            let status = PsLookupProcessByProcessId(pid, &mut p as *mut _ as _);
            if !NT_SUCCESS(status) {
                return Err(ProcessError {});
            }
        };
        let r = Self { raw: p, pid };
        Ok(r)
    }

    /// # Safety
    ///
    /// function will dereference phandle
    pub unsafe fn from_handle(h: PHANDLE) -> Result<Self, ProcessError> {
        let mut process: *mut c_void = core::ptr::null_mut();

        unsafe {
            let status = ObReferenceObjectByHandle(
                *h as _,
                THREAD_ALL_ACCESS,
                *PsProcessType,
                KernelMode as _,
                &mut process as *mut _ as _,
                core::ptr::null_mut(),
            );

            if !NT_SUCCESS(status) {
                println!("ObReferenceObjectByHandle process error:{}", status);
                return Err(ProcessError {});
            }
        }
        let pid = unsafe { PsGetProcessId(process as _) };

        let r = Self {
            raw: core::ptr::null_mut(),
            pid: pid,
        };

        Ok(r)
    }

    pub fn get_pid(self) -> *mut c_void {
        self.pid
    }

    pub fn get_eprocess(self) -> *mut c_void {
        self.raw
    }

    pub fn get_process_name(self) -> String {
        let name = unsafe { PsGetProcessImageFileName(self.raw as *mut _) };
        cstr_to_rust_str(name)
    }
}

impl Deref for Process {
    type Target = *mut c_void;

    fn deref(&self) -> &Self::Target {
        &self.raw
    }
}

impl DerefMut for Process {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.raw
    }
}

impl Drop for Process {
    fn drop(&mut self) {
        if !self.raw.is_null() {
            unsafe {
                ObfDereferenceObject(self.raw);
            }
        }
    }
}