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

DWORD GetPidByName(const wchar_t* exeName)
{
    if (!exeName)
        return 0;

    HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (snap == INVALID_HANDLE_VALUE)
        return 0;

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

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

    do
    {
        if (_wcsicmp(pe.szExeFile, exeName) == 0)
        {
            DWORD pid = pe.th32ProcessID;
            CloseHandle(snap);
            return pid;
        }
    } while (Process32Next(snap, &pe));

    CloseHandle(snap);
    return 0;
}

bool ProcessExistsByName(const wchar_t* exeName)
{
    return GetPidByName(exeName) != 0;
}

static bool BuildCommandLine(
    const wchar_t* exePath,
    const wchar_t* args,
    std::wstring& fullPath,
    std::wstring& commandLine)
{
    if (!exePath || !*exePath)
        return false;

    wchar_t abs[4096]{};
    DWORD n = GetFullPathNameW(exePath, _countof(abs), abs, nullptr);

    if (n == 0)
        wcscpy_s(abs, _countof(abs), exePath);

    fullPath = abs;

    if (GetFileAttributesW(fullPath.c_str()) == INVALID_FILE_ATTRIBUTES)
    {
        Log(L"[-] Executable not found: %s", fullPath.c_str());
        return false;
    }

    commandLine.clear();
    commandLine += L'"';
    commandLine += fullPath;
    commandLine += L'"';

    if (args && *args)
    {
        commandLine += L' ';
        commandLine += args;
    }

    return true;
}

static bool CreateProcessAsSystemImpl(
    HANDLE systemToken,
    const wchar_t* exePath,
    const wchar_t* args,
    bool interactive,
    bool suspended,
    NewProcInfo* out)
{
    if (out)
        *out = NewProcInfo();

    if (!systemToken || !exePath)
        return false;

    HANDLE tok = DuplicateSystemTokenFromHandle(systemToken);
    if (!tok)
    {
        Log(L"[-] Failed to duplicate SYSTEM token for process creation");
        return false;
    }

    EnablePrivilege(SE_ASSIGNPRIMARYTOKEN_NAME);
    EnablePrivilege(SE_INCREASE_QUOTA_NAME);

    if (interactive)
    {
        DWORD session = WTSGetActiveConsoleSessionId();

        if (session != 0xFFFFFFFF)
        {
            if (!SetTokenInformation(tok, TokenSessionId, &session, sizeof(session)))
            {
                Log(L"[!] SetTokenInformation(TokenSessionId) failed: %lu", GetLastError());
            }
        }
    }

    LPVOID env = nullptr;
    BOOL createdEnv = CreateEnvironmentBlock(&env, tok, FALSE);

    STARTUPINFOW si{};
    si.cb = sizeof(si);
    si.dwFlags = STARTF_USESHOWWINDOW;
    si.wShowWindow = interactive ? SW_SHOWDEFAULT : SW_HIDE;

    if (interactive)
        si.lpDesktop = const_cast<LPWSTR>(L"Winsta0\\Default");

    std::wstring fullPath;
    std::wstring cmd;

    if (!BuildCommandLine(exePath, args, fullPath, cmd))
    {
        if (createdEnv)
            DestroyEnvironmentBlock(env);

        CloseHandle(tok);
        return false;
    }

    PROCESS_INFORMATION pi{};

    DWORD creationFlags = 0;

    if (!interactive)
        creationFlags |= CREATE_NO_WINDOW;

    if (suspended)
        creationFlags |= CREATE_SUSPENDED;

    BOOL ok = CreateProcessAsUserW(
        tok,
        fullPath.c_str(),
        cmd.data(),
        nullptr,
        nullptr,
        FALSE,
        creationFlags,
        createdEnv ? env : nullptr,
        nullptr,
        &si,
        &pi);

    if (!ok)
    {
        Log(L"[-] CreateProcessAsUserW failed: %lu", GetLastError());
    }
    else
    {
        Log(L"[+] Created process: PID %lu", pi.dwProcessId);

        if (out)
        {
            out->hProcess = pi.hProcess;
            out->hThread = pi.hThread;
            out->processId = pi.dwProcessId;
        }
        else
        {
            if (pi.hThread)
                CloseHandle(pi.hThread);

            if (pi.hProcess)
                CloseHandle(pi.hProcess);
        }
    }

    if (createdEnv)
        DestroyEnvironmentBlock(env);

    CloseHandle(tok);

    return ok == TRUE;
}

bool StartProcessAsSystem(
    HANDLE systemToken,
    const wchar_t* exePath,
    const wchar_t* args,
    bool interactive)
{
    return CreateProcessAsSystemImpl(systemToken, exePath, args, interactive, false, nullptr);
}

bool CreateProcessSuspendedAsSystem(
    HANDLE systemToken,
    const wchar_t* exePath,
    const wchar_t* args,
    bool interactive,
    NewProcInfo* out)
{
    return CreateProcessAsSystemImpl(systemToken, exePath, args, interactive, true, out);
}