// trustinstaller_utils.cpp
#include "trustinstaller_utils.h"

void Elevate::EnablePrivilege(std::wstring privilegeName)
{
    HANDLE hToken;
    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &hToken))
        throw std::runtime_error("OpenProcessToken failed: " + std::to_string(GetLastError()));

    LUID luid;
    if (!LookupPrivilegeValueW(nullptr, privilegeName.c_str(), &luid))
    {
        CloseHandle(hToken);
        throw std::runtime_error("LookupPrivilegeValue failed: " + std::to_string(GetLastError()));
    }

    TOKEN_PRIVILEGES tp;
    tp.PrivilegeCount = 1;
    tp.Privileges[0].Luid = luid;
    tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
    if (!AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(TOKEN_PRIVILEGES), nullptr, nullptr))
    {
        CloseHandle(hToken);
        throw std::runtime_error("AdjustTokenPrivilege failed: " + std::to_string(GetLastError()));
    }

    CloseHandle(hToken);
}

DWORD Elevate::GetProcessIdByName(std::wstring  processName)
{
    HANDLE hSnapshot;
    if ((hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)) == INVALID_HANDLE_VALUE)
    {
        throw std::runtime_error("CreateToolhelp32Snapshot failed: " + std::to_string(GetLastError()));
    }

    DWORD pid = -1;
    PROCESSENTRY32W pe;
    ZeroMemory(&pe, sizeof(PROCESSENTRY32W));
    pe.dwSize = sizeof(PROCESSENTRY32W);
    if (Process32FirstW(hSnapshot, &pe))
    {
        while (Process32NextW(hSnapshot, &pe))
        {
            if (pe.szExeFile == processName)
            {
                pid = pe.th32ProcessID;
                break;
            }
        }
    }
    else
    {
        CloseHandle(hSnapshot);
        throw std::runtime_error("Process32First failed: " + std::to_string(GetLastError()));
    }

    if (pid == -1)
    {
        CloseHandle(hSnapshot);
        throw std::runtime_error("process not found: " + std::wstring_convert<std::codecvt_utf8<wchar_t>>().to_bytes(processName));
    }

    CloseHandle(hSnapshot);
    return pid;
}

void Elevate::ImpersonateSystem()
{
    auto systemPid = GetProcessIdByName(L"winlogon.exe");
    HANDLE hSystemProcess;
    if ((hSystemProcess = OpenProcess(
        PROCESS_DUP_HANDLE | PROCESS_QUERY_INFORMATION,
        FALSE,
        systemPid)) == nullptr)
    {
        throw std::runtime_error("OpenProcess failed (winlogon.exe): " + std::to_string(GetLastError()));
    }

    HANDLE hSystemToken;
    if (!OpenProcessToken(
        hSystemProcess,
        MAXIMUM_ALLOWED,
        &hSystemToken))
    {
        CloseHandle(hSystemProcess);
        throw std::runtime_error("OpenProcessToken failed (winlogon.exe): " + std::to_string(GetLastError()));
    }

    HANDLE hDupToken;
    SECURITY_ATTRIBUTES tokenAttributes;
    tokenAttributes.nLength = sizeof(SECURITY_ATTRIBUTES);
    tokenAttributes.lpSecurityDescriptor = nullptr;
    tokenAttributes.bInheritHandle = FALSE;
    if (!DuplicateTokenEx(
        hSystemToken,
        MAXIMUM_ALLOWED,
        &tokenAttributes,
        SecurityImpersonation,
        TokenImpersonation,
        &hDupToken))
    {
        CloseHandle(hSystemToken);
        throw std::runtime_error("DuplicateTokenEx failed (winlogon.exe): " + std::to_string(GetLastError()));
    }

    if (!ImpersonateLoggedOnUser(hDupToken))
    {
        CloseHandle(hDupToken);
        CloseHandle(hSystemToken);
        throw std::runtime_error("ImpersonateLoggedOnUser failed: " + std::to_string(GetLastError()));
    }

    CloseHandle(hDupToken);
    CloseHandle(hSystemToken);
}

int Elevate::StartTrustedInstallerService()
{
    SC_HANDLE hSCManager;
    if ((hSCManager = OpenSCManagerW(
        nullptr,
        SERVICES_ACTIVE_DATABASE,
        GENERIC_EXECUTE)) == nullptr)
    {
        throw std::runtime_error("OpenSCManager failed: " + std::to_string(GetLastError()));
    }

    SC_HANDLE hService;
    if ((hService = OpenServiceW(
        hSCManager,
        L"TrustedInstaller",
        GENERIC_READ | GENERIC_EXECUTE)) == nullptr)
    {
        CloseServiceHandle(hSCManager);
        throw std::runtime_error("OpenService failed: " + std::to_string(GetLastError()));
    }

    SERVICE_STATUS_PROCESS statusBuffer;
    DWORD bytesNeeded;
    while (QueryServiceStatusEx(
        hService,
        SC_STATUS_PROCESS_INFO,
        reinterpret_cast<LPBYTE>(&statusBuffer),
        sizeof(SERVICE_STATUS_PROCESS),
        &bytesNeeded))
    {
        if (statusBuffer.dwCurrentState == SERVICE_STOPPED)
        {
            if (!StartServiceW(hService, 0, nullptr))
            {
                CloseServiceHandle(hService);
                CloseServiceHandle(hSCManager);
                throw std::runtime_error("StartService failed: " + std::to_string(GetLastError()));
            }
        }
        if (statusBuffer.dwCurrentState == SERVICE_START_PENDING ||
            statusBuffer.dwCurrentState == SERVICE_STOP_PENDING)
        {
            Sleep(statusBuffer.dwWaitHint);
            continue;
        }
        if (statusBuffer.dwCurrentState == SERVICE_RUNNING)
        {
            CloseServiceHandle(hService);
            CloseServiceHandle(hSCManager);
            return statusBuffer.dwProcessId;
        }
    }
    CloseServiceHandle(hService);
    CloseServiceHandle(hSCManager);
    throw std::runtime_error("QueryServiceStatusEx failed: " + std::to_string(GetLastError()));
}

void Elevate::RunAsTrustedInstaller()
{
    auto pid = StartTrustedInstallerService();
    EnablePrivilege(SE_DEBUG_NAME);
    EnablePrivilege(SE_IMPERSONATE_NAME);
    Elevate::ImpersonateSystem();

    HANDLE hTIProcess;
    if ((hTIProcess = OpenProcess(
        PROCESS_DUP_HANDLE | PROCESS_QUERY_INFORMATION,
        FALSE,
        pid)) == nullptr)
    {
        throw std::runtime_error("OpenProcess failed (TrustedInstaller.exe): " + std::to_string(GetLastError()));
    }

    HANDLE hTIToken;
    if (!OpenProcessToken(
        hTIProcess,
        MAXIMUM_ALLOWED,
        &hTIToken))
    {
        CloseHandle(hTIProcess);
        throw std::runtime_error("OpenProcessToken failed (TrustedInstaller.exe): " + std::to_string(GetLastError()));
    }

    HANDLE hDupToken;
    SECURITY_ATTRIBUTES tokenAttributes;
    tokenAttributes.nLength = sizeof(SECURITY_ATTRIBUTES);
    tokenAttributes.lpSecurityDescriptor = nullptr;
    tokenAttributes.bInheritHandle = FALSE;
    if (!DuplicateTokenEx(
        hTIToken,
        MAXIMUM_ALLOWED,
        &tokenAttributes,
        SecurityImpersonation,
        TokenImpersonation,
        &hDupToken))
    {
        CloseHandle(hTIToken);
        throw std::runtime_error("DuplicateTokenEx failed (TrustedInstaller.exe): " + std::to_string(GetLastError()));
    }

    if (!ImpersonateLoggedOnUser(hTIToken)) {
        throw std::runtime_error("ImpersonateLoggedOnUser failed: " + std::to_string(GetLastError()));
    }
}