#include "bootmgr.h"
#include "efi_helper.h"
#include "common/io.h"
#include "common/windows/nt.h"

#include <ntstatus.h>
#include <windows.h>

const char* enablePrivilege(const wchar_t* privilege)
{
    FF_AUTO_CLOSE_FD HANDLE token = NULL;
    if (!NT_SUCCESS(NtOpenProcessToken(NtCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &token)))
        return "NtOpenProcessToken() failed";

    TOKEN_PRIVILEGES tp = {
        .PrivilegeCount = 1,
        .Privileges = {
            (LUID_AND_ATTRIBUTES) { .Attributes = SE_PRIVILEGE_ENABLED }
        },
    };
    if (!LookupPrivilegeValueW(NULL, privilege, &tp.Privileges[0].Luid))
        return "LookupPrivilegeValue() failed";

    NTSTATUS status = NtAdjustPrivilegesToken(token, false, &tp, sizeof(tp), NULL, NULL);
    if (!NT_SUCCESS(status))
        return "NtAdjustPrivilegesToken() failed";

    if (status == STATUS_NOT_ALL_ASSIGNED)
        return "The token does not have the specified privilege; try sudo please";

    return NULL;
}

const char* ffDetectBootmgr(FFBootmgrResult* result)
{
    const char* err = enablePrivilege(L"SeSystemEnvironmentPrivilege");
    if (err != NULL)
        return err;

    GUID efiGlobalGuid;
    if (!NT_SUCCESS(RtlGUIDFromString(&(UNICODE_STRING) RTL_CONSTANT_STRING(L"{" FF_EFI_GLOBAL_GUID L"}"), &efiGlobalGuid)))
        return "RtlGUIDFromString() failed";

    ULONG size = sizeof(result->order);
    if (!NT_SUCCESS(NtQuerySystemEnvironmentValueEx(&(UNICODE_STRING) RTL_CONSTANT_STRING(L"BootCurrent"), &efiGlobalGuid, &result->order, &size, NULL)))
        return "NtQuerySystemEnvironmentValueEx(BootCurrent) failed";
    if (size != sizeof(result->order))
        return "NtQuerySystemEnvironmentValueEx(BootCurrent) returned unexpected size";

    uint8_t buffer[2048];
    wchar_t key[9];
    swprintf(key, ARRAY_SIZE(key), L"Boot%04X", result->order);
    size = sizeof(buffer);
    if (!NT_SUCCESS(NtQuerySystemEnvironmentValueEx(&(UNICODE_STRING) RTL_CONSTANT_STRING(key), &efiGlobalGuid, buffer, &size, NULL)))
        return "NtQuerySystemEnvironmentValueEx(Boot####) failed";
    if (size < sizeof(FFEfiLoadOption) || size == ARRAY_SIZE(buffer))
        return "NtQuerySystemEnvironmentValueEx(Boot####) returned unexpected size";

    ffEfiFillLoadOption((FFEfiLoadOption *)buffer, result);

    SYSTEM_SECUREBOOT_INFORMATION ssi;
    if (NT_SUCCESS(NtQuerySystemInformation(SystemSecureBootInformation, &ssi, sizeof(ssi), NULL)))
        result->secureBoot = ssi.SecureBootEnabled;

    return NULL;
}
