moon-driver-utils 0.1.0

use core::ffi::c_void;

use wdk::println;
use wdk_sys::ntddk::{
    memcpy, ObGetFilterVersion, ObRegisterCallbacks, ObUnRegisterCallbacks, PsGetProcessId,
    PsRemoveCreateThreadNotifyRoutine, PsRemoveLoadImageNotifyRoutine,
    PsSetCreateThreadNotifyRoutine, PsSetLoadImageNotifyRoutine,
};
use wdk_sys::_OB_PREOP_CALLBACK_STATUS::OB_PREOP_SUCCESS;
use wdk_sys::{
    ntddk::PsSetCreateProcessNotifyRoutineEx, DRIVER_OBJECT, LIST_ENTRY64, UNICODE_STRING64,
};
use wdk_sys::{
    PsProcessType, NT_SUCCESS, OB_CALLBACK_REGISTRATION, OB_OPERATION_HANDLE_CREATE,
    OB_OPERATION_HANDLE_DUPLICATE, OB_OPERATION_REGISTRATION, OB_PREOP_CALLBACK_STATUS,
    POBJECT_TYPE, POB_POST_OPERATION_CALLBACK, POB_POST_OPERATION_INFORMATION,
    POB_PRE_OPERATION_CALLBACK, POB_PRE_OPERATION_INFORMATION, PVOID, UNICODE_STRING,
};

use crate::kernel_fucntion::{
    CreateImageNotifyRoutineExFn, CreateProcessNotifyRoutineExFn, CreateThreadNotifyRoutineExFn,
    PsGetProcessImageFileName,
};
use crate::string::{cstr_to_rust_str, str_to_unicode_string};

#[repr(C)]
struct LdrData {
    in_load_order_links: LIST_ENTRY64,
    in_memory_order_links: LIST_ENTRY64,
    in_initialization_order_links: LIST_ENTRY64,
    dll_base: *mut c_void,
    entry_point: *mut c_void,
    size_of_image: u32,
    padding0: [u8; 4],
    full_dll_name: UNICODE_STRING64,
    base_dll_name: UNICODE_STRING64,
    flags: u32,
}

#[derive(Default, Debug)]
pub struct CallbackUtil {}

impl CallbackUtil {
    pub fn new(driver_object: *mut DRIVER_OBJECT) -> Self {
        unsafe {
            let ldr: *mut LdrData = (*driver_object).DriverSection as _;
            let ldr = &mut (*ldr);
            ldr.flags = ldr.flags | 0x20;
        }

        Self {}
    }

    pub fn add_process_notify(&self, routine: CreateProcessNotifyRoutineExFn) {
        let status = unsafe { PsSetCreateProcessNotifyRoutineEx(Some(routine), 0) };
        if !NT_SUCCESS(status) {
            println!("add process notify fault:{}", status);
        }
    }

    pub fn remove_process_notify(&self, routine: CreateProcessNotifyRoutineExFn) {
        let status = unsafe { PsSetCreateProcessNotifyRoutineEx(Some(routine), 1) };
        if !NT_SUCCESS(status) {
            println!("remove process notify fault:{}", status);
        }
    }

    pub fn add_thread_notify(&self, routine: CreateThreadNotifyRoutineExFn) {
        let status = unsafe { PsSetCreateThreadNotifyRoutine(Some(routine)) };
        if !NT_SUCCESS(status) {
            println!("add thread notify fault:{}", status);
        }
    }

    pub fn remove_thread_notify(&self, routine: CreateThreadNotifyRoutineExFn) {
        let status = unsafe { PsRemoveCreateThreadNotifyRoutine(Some(routine)) };
        if !NT_SUCCESS(status) {
            println!("remove thread notify fault:{}", status);
        }
    }

    pub fn add_image_notify(&self, routine: CreateImageNotifyRoutineExFn) {
        let status = unsafe { PsSetLoadImageNotifyRoutine(Some(routine)) };
        if !NT_SUCCESS(status) {
            println!("add image notify fault:{}", status);
        }
    }

    pub fn remove_image_notify(&self, routine: CreateImageNotifyRoutineExFn) {
        let status = unsafe { PsRemoveLoadImageNotifyRoutine(Some(routine)) };
        if !NT_SUCCESS(status) {
            println!("remove image notify fault:{}", status);
        }
    }

    pub fn add_process_object_notify(
        &self,
        object_type: *mut POBJECT_TYPE,
        pre_call: POB_PRE_OPERATION_CALLBACK,
        post_call: POB_POST_OPERATION_CALLBACK,
    ) -> *mut c_void {
        let mut ob_reg: OB_CALLBACK_REGISTRATION = unsafe { core::mem::zeroed() };
        let mut op_reg: OB_OPERATION_REGISTRATION = unsafe { core::mem::zeroed() };

        ob_reg.Version = unsafe { ObGetFilterVersion() };
        ob_reg.OperationRegistrationCount = 1;
        ob_reg.RegistrationContext = core::ptr::null_mut();
        unsafe {
            let mut s = str_to_unicode_string("321124");
            memcpy(
                &mut ob_reg.Altitude as *mut _ as _,
                s.as_ptr() as _,
                core::mem::size_of::<UNICODE_STRING>() as _,
            );
        }
        ob_reg.OperationRegistration = &mut op_reg;

        op_reg.ObjectType = object_type;
        op_reg.Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
        op_reg.PreOperation = pre_call;
        op_reg.PostOperation = post_call;

        let mut handle: PVOID = core::ptr::null_mut();
        let status = unsafe { ObRegisterCallbacks(&mut ob_reg as _, &mut handle as *mut _) };
        if !NT_SUCCESS(status) {
            println!("ObRegisterCallbacks Error: {}", status);
            unsafe { ObUnRegisterCallbacks(handle) };
        }

        return handle;
    }
}

unsafe extern "C" fn pre_call(
    _registration_context: PVOID,
    operation_information: POB_PRE_OPERATION_INFORMATION,
) -> OB_PREOP_CALLBACK_STATUS {
    let operation_information = &mut (*operation_information);
    if operation_information.ObjectType != *PsProcessType {
        return OB_PREOP_SUCCESS;
    }

    let pid = PsGetProcessId(operation_information.Object as _);
    println!("pre_call pid:{:p}", pid);

    let name = PsGetProcessImageFileName(operation_information.Object as _);
    if pid as u64 == 0x111u64 {
        println!(
            "pre_call {} pid:{:p} eprocess:{:p}",
            cstr_to_rust_str(name),
            pid,
            operation_information.Object
        );
    }

    return OB_PREOP_SUCCESS;
}

unsafe extern "C" fn post_call(
    _registration_context: PVOID,
    _operation_information: POB_POST_OPERATION_INFORMATION,
) {
}

pub fn test_handle_callback(driver_obj: *mut DRIVER_OBJECT) {
    let call_back_util = CallbackUtil::new(driver_obj);
    let _handle = call_back_util.add_process_object_notify(
        unsafe { PsProcessType },
        Some(pre_call),
        Some(post_call),
    );
}