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)
}
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);
}
}
}
}