moon-driver-utils 0.1.0

use alloc::ffi::CString;
use core::ffi::c_void;

use moon_instructions::{read_gs_current_thread, read_msr};
use moon_struct::{
    inner::PKPROCESSOR_STATE,
    msr::msr_index::MSR_LSTAR,
    os::OS_INFO,
    pe::{ImageDosHeader, ImageExportDirectory, ImageFileHeader, ImageOptionalHeader64},
};
use wdk_sys::{
    ntddk::{
        memcpy, strcmp, IoGetCurrentProcess, KeStackAttachProcess, KeUnstackDetachProcess,
        MmGetSystemRoutineAddress, MmIsAddressValid, RtlCompareUnicodeString,
    },
    ACCESS_MASK, ACCESS_STATE, HANDLE, IMAGE_INFO, KAPC_STATE, KPROCESSOR_MODE, LIST_ENTRY64,
    NTSTATUS, OBJECT_ATTRIBUTES, PACCESS_STATE, PCLIENT_ID, PCONTEXT, PEPROCESS, PETHREAD,
    PHANDLE64, PIO_STACK_LOCATION, PIRP, POBJECT_ATTRIBUTES, POBJECT_TYPE, POOL_TYPE,
    PS_CREATE_NOTIFY_INFO, PULONG, PUNICODE_STRING, PVOID, SIZE_T, UCHAR, ULONG, ULONG64,
    UNICODE_STRING, _EXCEPTION_RECORD, _KPROCESS, _OBJECT_TYPE,
};

use crate::{process::get_process_by_name, string::str_to_unicode_string};

#[derive(Debug)]
pub struct GetFunctionAddressError {}

pub type ZwReadVirtualMemoryFn = unsafe extern "system" fn(
    process_handle: PHANDLE64,
    base_address: PVOID,
    buffer: PVOID,
    buffer_length: ULONG64,
    return_length: ULONG64,
) -> NTSTATUS;

pub type NtOpenProcessFn = unsafe extern "system" fn(
    process_handle: PHANDLE64,
    desired_access: ACCESS_MASK,
    object_attributes: POBJECT_ATTRIBUTES,
    client_id: PCLIENT_ID,
) -> NTSTATUS;

pub type PsSuspendThreadFn =
    unsafe extern "system" fn(thread: PETHREAD, previous_suspend_count: PULONG) -> NTSTATUS;

pub type PsResumeThreadFn =
    unsafe extern "system" fn(thread: PETHREAD, previous_suspend_count: PULONG) -> NTSTATUS;

pub type DbgkpSectionToFileHandleFn = unsafe extern "system" fn(section: PVOID) -> HANDLE;

pub type DbgkSendSystemDllMessagesFn =
    unsafe extern "system" fn(thread: PVOID, dbg_obj: PVOID, api_msg: PVOID) -> HANDLE;

pub type DbgkpWakeTargetFn = unsafe extern "system" fn(debug_event: PVOID);

pub type DbgkpPostModuleMessagesFn = unsafe extern "system" fn(
    process: *mut c_void,
    thread: *mut c_void,
    debug_object: *mut c_void,
) -> NTSTATUS;

pub type DbgkForwardExceptionFn = unsafe extern "system" fn(
    exception_record: *mut _EXCEPTION_RECORD,
    debug_exception: u8,
    second_chance: u8,
) -> u8;

pub type DbgkpSendApiMessageFn =
    unsafe extern "system" fn(s: u64, suspend_process: u8, dbg_km_msg: *mut c_void) -> NTSTATUS;

pub type DbgkClearProcessDebugObjectFn =
    unsafe extern "system" fn(process: *mut c_void, source_debug_object: *mut c_void) -> NTSTATUS;

pub type ZwSystemDebugControlFn = unsafe extern "system" fn(
    command: u32,
    input_buffer: PVOID,
    input_buffer_length: ULONG,
    output_buffer: PVOID,
    output_buffer_length: ULONG,
    return_length: PULONG,
) -> NTSTATUS;

pub type CreateProcessNotifyRoutineExFn = unsafe extern "C" fn(
    process: *mut _KPROCESS,
    process_id: HANDLE,
    create_info: *mut PS_CREATE_NOTIFY_INFO,
);

pub type CreateThreadNotifyRoutineExFn =
    unsafe extern "C" fn(process_id: HANDLE, thread_id: HANDLE, create: u8);

pub type CreateImageNotifyRoutineExFn = unsafe extern "C" fn(
    full_image_name: *mut UNICODE_STRING,
    process_id: HANDLE,
    image_info: *mut IMAGE_INFO,
);

pub type PspTerminateAllThreadsFn = unsafe extern "C" fn(
    a1: *mut c_void,
    a2: *mut c_void,
    a3: *mut c_void,
    a4: *mut c_void,
) -> NTSTATUS;

pub type DbgkCreateThreadFn = unsafe extern "C" fn(a1: *mut c_void) -> NTSTATUS;

pub type DbgkExitThreadFn = unsafe extern "C" fn(a1: *mut c_void) -> NTSTATUS;

pub type PspExitThreadFn = unsafe extern "C" fn(a1: *mut c_void) -> NTSTATUS;

pub type DbgkMapViewOfSectionFn =
    unsafe extern "C" fn(a1: *mut c_void, a2: *mut c_void, a3: *mut c_void) -> NTSTATUS;

pub type DbgkUnMapViewOfSectionFn =
    unsafe extern "C" fn(a1: *mut c_void, a2: *mut c_void) -> NTSTATUS;

pub fn get_current_thread() -> *mut c_void {
    read_gs_current_thread() as _
}

/// # Safety
///
/// derefernce to irp
pub unsafe fn io_get_current_irp_stack_location(irp: PIRP) -> PIO_STACK_LOCATION {
    unsafe {
        (*irp)
            .Tail
            .Overlay
            .__bindgen_anon_2
            .__bindgen_anon_1
            .CurrentStackLocation
    }
}

extern "C" {
    pub fn SeLocateProcessImageName(
        process: *mut c_void,
        image_file_name: *mut PUNICODE_STRING,
    ) -> NTSTATUS;
    pub fn PsGetProcessImageFileName(process: PEPROCESS) -> *mut UCHAR;
    pub fn ObGetObjectType(object: *mut c_void) -> *mut c_void;
    pub fn ObCreateObject(
        prebe_mode: KPROCESSOR_MODE,
        object_type: *mut _OBJECT_TYPE,
        object_attributes: *mut OBJECT_ATTRIBUTES,
        ownership_mode: KPROCESSOR_MODE,
        parse_context: *mut c_void,
        object_body_size: u32,
        paged_pool_charge: u32,
        no_paged_pool_charge: u32,
        object: *mut c_void,
    ) -> NTSTATUS;

    pub fn ObCreateObjectType(
        type_name: *mut UNICODE_STRING,
        object_type_init: *mut c_void,
        reserved: *mut c_void,
        object_type: *mut POBJECT_TYPE,
    ) -> NTSTATUS;

    pub fn ObInsertObject(
        object: PVOID,
        passed_access_state: PACCESS_STATE,
        desired_access: ACCESS_MASK,
        object_pointer_bias: u32,
        new_object: *mut PVOID,
        handle: PHANDLE64,
    ) -> NTSTATUS;
    pub fn KeSaveStateForHibernate(state: PKPROCESSOR_STATE);
    pub fn RtlRestoreContext(ContextRecord: PCONTEXT, ExceptionRecord: *mut _EXCEPTION_RECORD);
    pub fn ExAllocatePoolWithQuotaTag(
        PoolType: POOL_TYPE,
        NumberOfBytes: SIZE_T,
        Tag: ULONG,
    ) -> PVOID;

    pub fn ObDuplicateObject(
        source_process: *mut c_void,
        source_handle: HANDLE,
        target_process: *mut c_void,
        target_handle: *mut c_void,
        desired_access: ACCESS_MASK,
        handle_attributes: u32,
        options: u32,
        previous_mode: KPROCESSOR_MODE,
    ) -> NTSTATUS;

    pub fn ObReferenceObjectByName(
        object_name: *mut UNICODE_STRING,
        attributes: u32,
        passed_access_state: *mut ACCESS_STATE,
        desired_access: ACCESS_MASK,
        object_type: *mut c_void,
        access_mode: KPROCESSOR_MODE,
        parse_context: *mut c_void,
        object: *mut PVOID,
    ) -> NTSTATUS;

    pub fn ZwQueryDirectoryObject(
        directory_handle: HANDLE,
        buffer: *mut c_void,
        length: u32,
        return_single_entry: u8,
        restart_scan: u8,
        context: *mut u32,
        return_length: *mut u32,
    ) -> NTSTATUS;

    // pub fn NtOpenDirectoryObject(
    //     DirectoryHandle: *mut HANDLE,
    //     DesiredAccess: ACCESS_MASK,
    //     ObjectAttributes: POBJECT_ATTRIBUTES,
    // ) -> NTSTATUS;
}

#[repr(C)]
#[derive(Debug)]
pub struct SystemServiceTable {
    pub service_table_base: *mut u32,
    service_counter_table_base: *mut core::ffi::c_void,
    number_of_services: usize,
    param_table_base: *mut core::ffi::c_void,
}

#[allow(unused)]
pub fn get_ssdt_address() -> *mut SystemServiceTable {
    let start_search_address = read_msr(MSR_LSTAR);
    let end_search_address = start_search_address + 0x500;

    let mut p = 0u64;

    let mut b1 = 0u8;
    let mut b2 = 0u8;
    let mut b3 = 0u8;

    for i in start_search_address..end_search_address {
        unsafe {
            if MmIsAddressValid(i as _) != 0
                && MmIsAddressValid((i + 1) as _) != 0
                && MmIsAddressValid((i + 2) as _) != 0
            {
                b1 = *(i as *mut u8);
                b2 = *((i + 1) as *mut u8);
                b3 = *((i + 2) as *mut u8);
                if b1 == 0x4c && b2 == 0x8d && b3 == 0x15 {
                    let mut temp = 0u32;
                    memcpy(&mut temp as *mut u32 as _, (i + 3) as _, 4 as _);
                    p = temp as u64 + i + 7u64;
                }
            }
        }
    }

    p as _
}

pub fn get_ldr_module_base_by_name(name: &str) -> *mut c_void {
    let mut name = str_to_unicode_string(name);
    let cprocess: *mut c_void = unsafe { IoGetCurrentProcess() } as _;

    unsafe {
        let peb = *(cprocess.add(OS_INFO.offset.eprocess_peb) as *mut u64) as *mut c_void;
        if peb.is_null() {
            // driver
            return core::ptr::null_mut();
        }

        let ldr = *(peb.add(OS_INFO.offset.peb_ldr) as *mut u64) as *mut c_void;

        let list_head = ldr.add(OS_INFO.offset.ldr_in_load_order_module_list) as *mut LIST_ENTRY64;

        let mut c = (*list_head).Flink as *mut LIST_ENTRY64;

        let mut count = 0;
        loop {
            if c == list_head || count >= 0x1f4 {
                break;
            }

            if RtlCompareUnicodeString(
                ((c as u64) + (OS_INFO.offset.ldre_base_dll_name as u64)) as _,
                name.as_ptr(),
                0,
            ) == 0
            {
                return *((c as u64 + (OS_INFO.offset.ldre_dll_base as u64)) as *mut u64) as _;
            }

            c = (*c).Flink as _;
            count += 1;
        }
    }

    core::ptr::null_mut()
}

/// # Safety
///
/// dereference module point
pub unsafe fn get_module_export_address(module: *mut c_void, name: &str) -> *mut c_void {
    let dos_header = module as *mut ImageDosHeader;
    if dos_header.is_null() {
        return core::ptr::null_mut();
    }
    unsafe {
        let option_header: *mut ImageOptionalHeader64 = (module as u64
            + (*dos_header).e_lfanew as u64
            + core::mem::size_of::<u32>() as u64
            + core::mem::size_of::<ImageFileHeader>() as u64)
            as _;
        if option_header.is_null() {
            return core::ptr::null_mut();
        }

        let export_table: *mut ImageExportDirectory =
            (module as u64 + (*option_header).data_directory[0].virtual_address as u64) as _;

        if export_table.is_null() {
            return core::ptr::null_mut();
        }

        let size = (*option_header).data_directory[0].size;

        let name_table_base = (module as u64 + (*export_table).address_of_names as u64) as *mut u32;

        // 如果通过id，就可用直接取name_ordinal_table_base[id]
        let name_ordinal_table_base =
            (module as u64 + (*export_table).address_of_name_ordinals as u64) as *mut u16;

        let address_table_base =
            (module as u64 + (*export_table).address_of_functions as u64) as *mut u32;

        let mut low = 0u32;
        let mut middle = 0u32;
        let mut hight = (*export_table).number_of_names - 1;

        let cname = CString::new(name).unwrap();

        // 二分找导出函数
        while hight >= low {
            middle = (low + hight) >> 1;

            let s1 = cname.as_ptr();
            let s2 = module.add(*(name_table_base.add(middle as _)) as _);

            let result = strcmp(s1 as _, s2 as _);

            if result < 0 {
                hight = middle - 1;
            } else if result > 0 {
                low = middle + 1;
            } else {
                break;
            }
        }

        if hight < low {
            return core::ptr::null_mut();
        }

        let ordinal_number = *(name_ordinal_table_base.add(middle as _));
        if ordinal_number > (*export_table).number_of_functions as _ {
            return core::ptr::null_mut();
        }

        let function_address = module.add(*(address_table_base.add(ordinal_number as _)) as _);
        let end = export_table as u64 + size as u64;
        if function_address as u64 > export_table as u64 && (function_address as u64) < end {
            return core::ptr::null_mut();
        }

        function_address as _
    }
}

pub fn get_ntdll_function_id(name: &str) -> u32 {
    let Ok(process) = get_process_by_name("smss.exe") else {
        return 0;
    };

    let mut result = u32::MAX;

    unsafe {
        let mut apc_state = KAPC_STATE::default();
        KeStackAttachProcess(process, &mut apc_state as _);

        let ntdll = get_ldr_module_base_by_name("ntdll.dll");
        if !ntdll.is_null() {
            let func = get_module_export_address(ntdll, name);
            if !func.is_null() {
                result = *((func.add(4)) as *mut u16) as _;
            }
        }

        KeUnstackDetachProcess(&mut apc_state as _);
    }

    result
}

pub fn get_ssdt_address_by_id(id: u32) -> *mut c_void {
    let ssdt = get_ssdt_address();
    if ssdt.is_null() {
        return core::ptr::null_mut();
    }

    unsafe {
        let base = (*ssdt).service_table_base;
        if base.is_null() {
            return core::ptr::null_mut();
        }

        let mut temp = *base.add(id as _);
        temp >>= 4;
        (temp as u64 + base as u64) as _
    }
}

pub fn get_ssdt_function_address_by_name(
    name: &str,
) -> Result<*mut c_void, GetFunctionAddressError> {
    let id = get_ntdll_function_id(name);

    if id != u32::MAX {
        return Ok(get_ssdt_address_by_id(id));
    }

    Err(GetFunctionAddressError {})
}

pub fn get_zw_function_address_by_id(nt_function_id: u64) -> *mut c_void {
    let mut zw_close_us = str_to_unicode_string("ZwClose");
    let zw_close_adderss = unsafe { MmGetSystemRoutineAddress(zw_close_us.as_ptr()) };

    let mut zw_close_index: u32 = 0;
    let zw_function_len: u32 = 32;

    unsafe {
        memcpy(
            &mut zw_close_index as *mut _ as *mut _,
            zw_close_adderss.add(21),
            4,
        )
    };

    let f0_addr = zw_close_adderss as u64 - (zw_close_index * zw_function_len) as u64;
    (f0_addr + nt_function_id * zw_function_len as u64) as _
}

pub fn get_kernel_export_symbol_address(
    name: &str,
) -> Result<*mut c_void, GetFunctionAddressError> {
    let mut function_name = str_to_unicode_string(name);
    unsafe {
        let r = MmGetSystemRoutineAddress(function_name.as_ptr() as _);
        if MmIsAddressValid(r) != 0 {
            return Ok(r);
        }

        Err(GetFunctionAddressError {})
    }
}