kdstinker/kdstinker/driver_util.cpp

#include "driver_util.h"

namespace driver_util
{
    void* get_driver_base(const char* module_name)
    {
        ULONG bytes{};
        NTSTATUS status = ZwQuerySystemInformation(
            SystemModuleInformation,
            NULL,
            bytes,
            &bytes
        );

        if (!bytes)
            return NULL;
        PRTL_PROCESS_MODULES modules =
            (PRTL_PROCESS_MODULES)ExAllocatePool(NonPagedPool, bytes);

        if (modules)
        {
            status = ZwQuerySystemInformation(
                SystemModuleInformation,
                modules,
                bytes,
                &bytes
            );

            if (!NT_SUCCESS(status))
            {
                ExFreePool(modules);
                return NULL;
            }

            PRTL_PROCESS_MODULE_INFORMATION module = modules->Modules;
            PVOID module_base{}, module_size{};
            for (ULONG i = 0; i < modules->NumberOfModules; i++)
            {
                if (strcmp(reinterpret_cast<char*>(module[i].FullPathName + module[i].OffsetToFileName), module_name) == 0)
                {
                    module_base = module[i].ImageBase;
                    module_size = (PVOID)module[i].ImageSize;
                    break;
                }
            }
            ExFreePool(modules);
            return module_base;
        }
        return NULL;
    }

    void* get_kmode_export(const char* mod_name, const char* proc_name)
    {
        if (!mod_name || !proc_name)
            return NULL;

        void* result = get_driver_base(mod_name);
        if (!result)
            return NULL;
        return RtlFindExportedRoutineByName(result, proc_name);
    }

    PIMAGE_FILE_HEADER get_file_header(void* base_addr)
    {
        if (!base_addr || *(short*)base_addr != 0x5A4D)
            return NULL;

        PIMAGE_DOS_HEADER dos_headers =
            reinterpret_cast<PIMAGE_DOS_HEADER>(base_addr);

        PIMAGE_NT_HEADERS nt_headers =
            reinterpret_cast<PIMAGE_NT_HEADERS>(
                reinterpret_cast<DWORD_PTR>(base_addr) + dos_headers->e_lfanew);

        return &nt_headers->FileHeader;
    }

    void* iat_hook(void* base_addr, const char* import, void* func_addr)
    {
        if (!base_addr || *(short*)base_addr != 0x5A4D || !import || !func_addr)
            return NULL;

        PIMAGE_DOS_HEADER dos_headers =
            reinterpret_cast<PIMAGE_DOS_HEADER>(base_addr);

        PIMAGE_NT_HEADERS nt_headers =
            reinterpret_cast<PIMAGE_NT_HEADERS>(
                reinterpret_cast<DWORD_PTR>(base_addr) + dos_headers->e_lfanew);

        IMAGE_DATA_DIRECTORY import_dir =
            nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];

        PIMAGE_IMPORT_DESCRIPTOR import_des =
            reinterpret_cast<PIMAGE_IMPORT_DESCRIPTOR>(import_dir.VirtualAddress + (DWORD_PTR)base_addr);

        LPCSTR lib_name = NULL;
        PVOID result = NULL;
        PIMAGE_IMPORT_BY_NAME func_name = NULL;

        if (!import_des)
            return NULL;

        while (import_des->Name != NULL)
        {
            lib_name = (LPCSTR)import_des->Name + (DWORD_PTR)base_addr;

            if (get_driver_base(lib_name))
            {
                PIMAGE_THUNK_DATA org_first_thunk = NULL, first_thunk = NULL;
                org_first_thunk = (PIMAGE_THUNK_DATA)((DWORD_PTR)base_addr + import_des->OriginalFirstThunk);
                first_thunk = (PIMAGE_THUNK_DATA)((DWORD_PTR)base_addr + import_des->FirstThunk);
                while (org_first_thunk->u1.AddressOfData != NULL)
                {
                    func_name = (PIMAGE_IMPORT_BY_NAME)((DWORD_PTR)base_addr + org_first_thunk->u1.AddressOfData);
                    if (strcmp(func_name->Name, import) == 0)
                    {
                        // save old function pointer
                        result = reinterpret_cast<PVOID>(first_thunk->u1.Function);

                        {
                            //
                            // disable write protection
                            //
                            auto cr0 = __readcr0();
                            cr0 &= 0xfffffffffffeffff;
                            __writecr0(cr0);
                            _disable();
                        }

                        // swap address
                        first_thunk->u1.Function = reinterpret_cast<ULONG64>(func_addr);
                        
                        {
                            //
                            // enable write protection
                            //
                            auto cr0 = __readcr0();
                            cr0 |= 0x10000;
                            _enable();
                            __writecr0(cr0);
                        }
                        return result;
                    }
                    ++org_first_thunk;
                    ++first_thunk;
                }
            }
            ++import_des;
        }
        return NULL;
    }

    void mem_dump(void* base_addr, unsigned len)
    {
        if (!base_addr || !len)
            return;

        HANDLE             h_file;
        UNICODE_STRING     name;
        OBJECT_ATTRIBUTES  attr;
        IO_STATUS_BLOCK    status_block;
        LARGE_INTEGER      offset{ NULL };

        RtlInitUnicodeString(&name, L"\\DosDevices\\C:\\dump.bin");
        InitializeObjectAttributes(&attr, &name,
            OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE,
            NULL, NULL
        );

        auto status = ZwCreateFile(
            &h_file,
            GENERIC_WRITE,
            &attr,
            &status_block,
            NULL,
            FILE_ATTRIBUTE_NORMAL,
            NULL,
            FILE_OVERWRITE_IF,
            FILE_SYNCHRONOUS_IO_NONALERT,
            NULL,
            NULL
        );

        status = ZwWriteFile(
            h_file,
            NULL,
            NULL,
            NULL,
            &status_block,
            base_addr,
            len,
            &offset,
            NULL
        );
        ZwClose(h_file);
    }
}