moon-driver-utils 0.1.0

use core::ffi::c_void;

use moon_struct::inner::{DIRECTORY_BASIC_INFORMATION, KLDR_DATA_TABLE_ENTRY};
use wdk::println;
use wdk_sys::{
    ntddk::ZwOpenDirectoryObject, DIRECTORY_QUERY, DRIVER_OBJECT, HANDLE, NT_SUCCESS,
    OBJECT_ATTRIBUTES, OBJ_CASE_INSENSITIVE, STATUS_BUFFER_TOO_SMALL, STATUS_MORE_ENTRIES,
    _DRIVER_OBJECT, _MODE::KernelMode,
};

use crate::{
    init_obj_attr,
    kernel_fucntion::{
        get_kernel_export_symbol_address, ObReferenceObjectByName, ZwQueryDirectoryObject,
    },
    kernel_struct::{IoDriverObjectType, NotFind},
    memory::pp::PP,
    string::{str_to_unicode_string, unicode_string_to_string},
};

pub fn find_kernel_base() -> Result<*mut c_void, NotFind> {
    let addr = get_kernel_export_symbol_address("PsLoadedModuleList");
    if addr.is_err() {
        return Err(NotFind {});
    }

    let addr = addr.unwrap();

    unsafe {
        let p = *(addr as *mut u64);

        let p = &(*(p as *mut KLDR_DATA_TABLE_ENTRY));

        return Ok(p.DllBase);
    }
}

pub fn find_driver_obj_no_device() -> Result<*mut _DRIVER_OBJECT, NotFind> {
    let mut h: HANDLE = core::ptr::null_mut();
    let mut obj_attr: OBJECT_ATTRIBUTES = unsafe { core::mem::zeroed() };
    let mut name = str_to_unicode_string("\\Driver");
    init_obj_attr(&mut obj_attr, name.as_ptr());

    let r = unsafe { ZwOpenDirectoryObject(&mut h as _, DIRECTORY_QUERY, &mut obj_attr as *mut _) };
    if !NT_SUCCESS(r) {
        println!("ZwOpenDirectoryObject UNSUCCESS:{}", r);
        return Err(NotFind {});
    }

    let mut size = 0x1000;
    let mut input_buffer;
    loop {
        let input_buffer2 = PP::<u8>::new_bytes(size);
        if input_buffer2.is_err() {
            println!("if input_buffer.is_err()");
            return Err(NotFind {});
        }

        input_buffer = input_buffer2.unwrap();

        let mut context: u32 = 0;
        let mut len: u32 = 0;
        let r = unsafe {
            ZwQueryDirectoryObject(
                h,
                input_buffer.as_ptr() as _,
                0x1000,
                0,
                1,
                &mut context as _,
                &mut len as _,
            )
        };
        if !NT_SUCCESS(r) {
            if r == STATUS_MORE_ENTRIES || r == STATUS_BUFFER_TOO_SMALL {
                size = size * 2;
                continue;
            }

            println!("ZwQueryDirectoryObject UNSUCCSS:{}", r);
            return Err(NotFind {});
        }

        break;
    }

    let mut p: *mut DIRECTORY_BASIC_INFORMATION = input_buffer.as_ptr().cast();

    loop {
        let it = unsafe { p.as_mut().unwrap() };
        if it.object_name.Length == 0 || it.object_type_name.Length == 0 {
            break;
        }

        let object_name = unicode_string_to_string(&it.object_name);
        // println!("item:{}", object_name);

        let mut name = str_to_unicode_string(&alloc::format!(
            "\\Driver\\{}",
            unicode_string_to_string(&it.object_name)
        ));
        let mut obj: *mut DRIVER_OBJECT = core::ptr::null_mut();
        let r = unsafe {
            ObReferenceObjectByName(
                name.as_ptr(),
                OBJ_CASE_INSENSITIVE,
                core::ptr::null_mut(),
                0,
                *(IoDriverObjectType as *mut u64) as _,
                KernelMode as _,
                core::ptr::null_mut(),
                &mut obj as *mut _ as _,
            )
        };

        if !NT_SUCCESS(r) {
            println!("fault: status:{:X}", r);
        } else {
            unsafe {
                let device_obj = obj.as_mut().unwrap().DeviceObject;
                if device_obj.is_null() {
                    // println!("obj:{:p},devcie:{:p}", obj, device_obj);
                    if object_name == "Win32k" {
                        return Ok(p as _);
                    }
                }
            }
        }

        p = unsafe { p.add(1) };
    }

    return Err(NotFind {});
}