#include <ntdll.h>
#include <pe.h>
#include <msg.h>

#include <ace.h>

static void _dll_notification(ULONG reason, const PLDR_DLL_NOTIFICATION_DATA data, void *context) {
    if (reason != 1) { // 1 - attach
        return;
    }

    // context should be set to the target module name
    wchar_t *targetModuleName = (wchar_t*)context;

    if (wcsicmp(targetModuleName, data->Loaded.BaseDllName->Buffer) != 0) {
        return;
    }

    // Replace entry point with a stub
    void *entryPoint = pe_find_entry_point(data->Loaded.DllBase);

    const char ENTRY_POINT_STUB[] = {
        0xB8, 0x01, 0x00, 0x00, 0x00,   // mov eax, 1
        0xC3                            // ret
    };

    DWORD oldProtect;
    VirtualProtect(entryPoint, sizeof(ENTRY_POINT_STUB), PAGE_EXECUTE_READWRITE, &oldProtect);

    memcpy(entryPoint, ENTRY_POINT_STUB, sizeof(ENTRY_POINT_STUB));

    VirtualProtect(entryPoint, sizeof(ENTRY_POINT_STUB), oldProtect, &oldProtect);
    
}

static void _create_driver_file(const char *path) {
    // They only report presence
    HANDLE file = CreateFileA(path, GENERIC_WRITE, FILE_SHARE_READ, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);

    if (file == INVALID_HANDLE_VALUE) {
        msg_err_a("Could not create driver file: %s", path);
    }

    CloseHandle(file);
}

void ace_fake_driver_files() {
    _create_driver_file("ACE-BASE.sys");

    // Just in case
    _create_driver_file("C:\\windows\\system32\\drivers\\ACE-BASE.sys");
}

static HMODULE _load_module_patched(wchar_t *path) {
    // Get filename from the path
    wchar_t *name = wcsrchr(path, '\\');
    name = name ? name + 1 : path;

    void *cookie;
    LdrRegisterDllNotification(0, &_dll_notification, name, &cookie);

    HMODULE module = LoadLibraryW(path);
    if (!module) {
        msg_err_w(L"Could not load module: %ls", path);
    }

    // LoadLibraryW is synchronous; the notification function has already finished executing
    LdrUnregisterDllNotification(cookie);

    return module;
}

HMODULE ace_load_base_module(const char *exeName) {
    wchar_t baseModuleName[MAX_PATH];
    swprintf(baseModuleName, MAX_PATH, L"%sBase.dll", exeName);
    wcslwr(baseModuleName);

    return _load_module_patched(baseModuleName);
}

HMODULE ace_load_driver_module() {
    return _load_module_patched(L"AntiCheatExpert\\InGame\\x64\\ACE-DRV64.dll");
}