#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, lowercase
    wchar_t *targetModuleName = (wchar_t*)context;

    wchar_t lwModuleName[MAX_PATH];
    wcscpy(lwModuleName, data->Loaded.BaseDllName->Buffer);
    _wcslwr(lwModuleName);

    if (wcscmp(targetModuleName, lwModuleName) == 0) {
        // 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);
    }
}

void ace_fake_driver_files() {
    // They only report presence
    const char *wdDriverPath = "ACE-BASE.sys";
    const char *s32DriverPath = "C:\\windows\\system32\\drivers\\ACE-BASE.sys";

    HANDLE wdDriverFile = CreateFileA(wdDriverPath, GENERIC_WRITE, FILE_SHARE_READ, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    if (!wdDriverFile) {
        msg_err_a("Could not create driver file: %s", wdDriverPath);
    }

    // Just in case
    HANDLE s32DriverFile = CreateFileA(s32DriverPath, GENERIC_WRITE, FILE_SHARE_READ, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    if (!s32DriverFile) {
        msg_err_a("Could not create driver file: %s", s32DriverPath);
    }

    CloseHandle(wdDriverFile);
    CloseHandle(s32DriverFile);
}

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

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

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

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

    return baseModule;
}

HMODULE ace_load_driver_module() {
    const char *driverModulePath = "AntiCheatExpert/InGame/x64/ACE-DRV64.dll";

    void *cookie;
    LdrRegisterDllNotification(0, &_dll_notification, L"ace-drv64.dll", &cookie);

    HMODULE driverModule = LoadLibraryA(driverModulePath);
    if (!driverModule) {
        msg_err_a("Could not load driver module: %s", driverModulePath);
    }

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

    return driverModule;
}