#include "Common.h"
#include "KillSafety.h"
#include "DriverClient.h"
#include "Logger.h"

#include <algorithm>
#include <cwctype>
#include <set>
#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 BaseNameFromPath(const std::wstring& path)
{
    size_t slash = path.find_last_of(L"\\/");

    if (slash == std::wstring::npos)
        return path;

    return path.substr(slash + 1);
}

bool IsProtectedImage(const std::wstring& imageName)
{
    static const wchar_t* protectedImages[] =
    {
        // Kernel / session manager / critical subsystems. Terminate any of
        // these and the machine bugchecks immediately.
        L"system",
        L"registry",
        L"memcompression",
        L"smss.exe",
        L"csrss.exe",
        L"wininit.exe",
        L"winlogon.exe",
        L"services.exe",
        L"lsass.exe",
        L"svchost.exe",
        L"fontdrvhost.exe",
        L"dwm.exe",
        L"wudfhost.exe",
        L"sihost.exe",
        L"ctfmon.exe",
        L"explorer.exe",
        L"taskhostw.exe",
        L"logonui.exe",
        L"userinit.exe",
        L"wininit.exe",
        L"spoolsv.exe",
        L"audiodg.exe",
        L"conhost.exe",
        L"sshd.exe",
    };

    const std::wstring lower = Lower(imageName);

    for (const wchar_t* p : protectedImages)
    {
        if (lower == Lower(p))
            return true;
    }

    return false;
}

// Collect our own pid and every ancestor up to the root.
static std::set<DWORD> SelfAndAncestors()
{
    std::set<DWORD> result;

    HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if (snap == INVALID_HANDLE_VALUE)
    {
        result.insert(GetCurrentProcessId());
        return result;
    }

    std::vector<std::pair<DWORD, DWORD>> links; // pid -> parent pid
    PROCESSENTRY32W pe{};
    pe.dwSize = sizeof(pe);

    if (Process32First(snap, &pe))
    {
        do
        {
            links.push_back({ pe.th32ProcessID, pe.th32ParentProcessID });
        } while (Process32Next(snap, &pe));
    }

    CloseHandle(snap);

    DWORD cur = GetCurrentProcessId();

    for (int depth = 0; depth < 16 && cur != 0; ++depth)
    {
        if (!result.insert(cur).second)
            break;

        DWORD parent = 0;

        for (const auto& link : links)
        {
            if (link.first == cur)
            {
                parent = link.second;
                break;
            }
        }

        cur = parent;
    }

    return result;
}

static std::vector<DWORD> PidsForImage(const std::wstring& imageName)
{
    std::vector<DWORD> pids;

    HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if (snap == INVALID_HANDLE_VALUE)
        return pids;

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

    if (Process32First(snap, &pe))
    {
        do
        {
            if (_wcsicmp(pe.szExeFile, imageName.c_str()) == 0)
                pids.push_back(pe.th32ProcessID);
        } while (Process32Next(snap, &pe));
    }

    CloseHandle(snap);
    return pids;
}

bool KillReport::Refused(KillOutcome o) const
{
    return o == KillOutcome::Protected ||
           o == KillOutcome::SelfOrAncestor ||
           o == KillOutcome::IdentityMismatch;
}

std::wstring KillReport::Format() const
{
    std::wostringstream os;

    os << (dryRun ? L"DRY-RUN" : L"EXECUTED")
       << L" killed=" << killed
       << L" candidates=" << attempts.size();

    const size_t maxLines = 30;
    size_t shown = 0;

    for (const auto& a : attempts)
    {
        if (shown >= maxLines)
        {
            os << L"\r\n  ... " << (attempts.size() - shown) << L" more";
            break;
        }

        ++shown;

        const wchar_t* state = L"?";

        switch (a.outcome)
        {
        case KillOutcome::WouldKill:       state = L"would-kill"; break;
        case KillOutcome::Killed:          state = L"killed"; break;
        case KillOutcome::Protected:       state = L"SKIP protected"; break;
        case KillOutcome::SelfOrAncestor:  state = L"SKIP own-tree"; break;
        case KillOutcome::IdentityMismatch:state = L"SKIP identity"; break;
        case KillOutcome::OpenFailed:      state = L"skip open-failed"; break;
        case KillOutcome::AccessDenied:    state = L"SKIP access-denied(needs SYSTEM or PPL)"; break;
        case KillOutcome::DriverFailed:    state = L"driver-failed"; break;
        }

        os << L"\r\n  " << a.pid << L" " << state;

        if (!a.imagePath.empty())
            os << L" " << a.imagePath;
    }

    return os.str();
}

void KillProcessesSafe(
    const std::vector<std::wstring>& imageNames,
    bool dryRun,
    DWORD maxKills,
    KillReport& report,
    bool allowBlindDriver)
{
    report = KillReport();
    report.dryRun = dryRun;

    if (maxKills == 0)
        maxKills = 16;

    const std::set<DWORD> own = SelfAndAncestors();

    std::set<DWORD> seen;
    DWORD kills = 0;

    for (const auto& requested : imageNames)
    {
        if (kills >= maxKills)
            break;

        if (requested.empty())
            continue;

        const std::wstring requestedBase = BaseNameFromPath(requested);

        // Refuse before enumerating: a protected image is never a candidate.
        if (IsProtectedImage(requestedBase))
        {
            KillAttempt blocked;
            blocked.outcome = KillOutcome::Protected;
            blocked.imagePath = requestedBase;
            report.attempts.push_back(blocked);
            continue;
        }

        for (DWORD pid : PidsForImage(requestedBase))
        {
            if (kills >= maxKills)
                break;

            if (pid == 0 || pid == 4 || pid <= 4)
                continue;

            if (!seen.insert(pid).second)
                continue;

            KillAttempt attempt;
            attempt.pid = pid;

            if (own.count(pid))
            {
                attempt.outcome = KillOutcome::SelfOrAncestor;
                report.attempts.push_back(attempt);
                continue;
            }

            // Hold the handle across the IOCTL: the process object stays alive,
            // so the PID cannot be recycled into something critical mid-kill.
            HANDLE proc = OpenProcess(
                PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_TERMINATE,
                FALSE,
                pid);

            if (!proc)
            {
                // User-mode terminate is refused for PPL / SYSTEM images. The
                // driver path does not need PROCESS_TERMINATE, so verify identity
                // through a query handle and hand the PID to the kernel.
                HANDLE queryOnly = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pid);

                if (!queryOnly)
                {
                    // Nothing openable (SYSTEM image while we are not elevated).
                    // Only the kernel can still act, and only when the caller
                    // explicitly allowed it: re-check the PID against a fresh
                    // snapshot so a recycled PID cannot slip through.
                    bool stillMatches = false;

                    for (DWORD again : PidsForImage(requestedBase))
                    {
                        if (again == pid)
                        {
                            stillMatches = true;
                            break;
                        }
                    }

                    if (allowBlindDriver && stillMatches && !IsProtectedImage(requestedBase))
                    {
                        if (dryRun)
                        {
                            attempt.outcome = KillOutcome::WouldKill;
                            attempt.imagePath += L" (driver)";
                            report.attempts.push_back(attempt);
                            continue;
                        }

                        const bool viaDriver = TerminateProcessViaDriver(pid);

                        if (viaDriver)
                        {
                            attempt.outcome = KillOutcome::Killed;
                            ++kills;
                            report.killed = (int)kills;
                        }
                        else
                        {
                            attempt.outcome = KillOutcome::DriverFailed;
                        }

                        report.attempts.push_back(attempt);
                        Sleep(150 + (GetTickCount() % 200));
                        continue;
                    }

                    attempt.outcome =
                        (GetLastError() == ERROR_ACCESS_DENIED)
                        ? KillOutcome::AccessDenied
                        : KillOutcome::OpenFailed;

                    report.attempts.push_back(attempt);
                    continue;
                }

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

                if (!QueryFullProcessImageNameW(queryOnly, 0, fullPath, &len))
                {
                    CloseHandle(queryOnly);
                    attempt.outcome = KillOutcome::OpenFailed;
                    report.attempts.push_back(attempt);
                    continue;
                }

                attempt.imagePath.assign(fullPath, len);

                const std::wstring actualBase = BaseNameFromPath(attempt.imagePath);

                if (_wcsicmp(actualBase.c_str(), requestedBase.c_str()) != 0)
                {
                    CloseHandle(queryOnly);
                    attempt.outcome = KillOutcome::IdentityMismatch;
                    report.attempts.push_back(attempt);
                    continue;
                }

                if (IsProtectedImage(actualBase))
                {
                    CloseHandle(queryOnly);
                    attempt.outcome = KillOutcome::Protected;
                    report.attempts.push_back(attempt);
                    continue;
                }

                if (dryRun)
                {
                    CloseHandle(queryOnly);
                    attempt.outcome = KillOutcome::WouldKill;
                    report.attempts.push_back(attempt);
                    continue;
                }

                const bool viaDriver = TerminateProcessViaDriver(pid);

                CloseHandle(queryOnly);

                if (viaDriver)
                {
                    attempt.outcome = KillOutcome::Killed;
                    ++kills;
                    report.killed = (int)kills;
                }
                else
                {
                    attempt.outcome = KillOutcome::DriverFailed;
                }

                report.attempts.push_back(attempt);

                Sleep(150 + (GetTickCount() % 200));
                continue;
            }

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

            if (!QueryFullProcessImageNameW(proc, 0, fullPath, &len))
            {
                CloseHandle(proc);
                attempt.outcome = KillOutcome::OpenFailed;
                report.attempts.push_back(attempt);
                continue;
            }

            attempt.imagePath.assign(fullPath, len);

            const std::wstring actualBase = BaseNameFromPath(attempt.imagePath);

            // Identity must match what we asked to kill, otherwise it is a
            // different image that happened to share a PID or a name variant.
            if (_wcsicmp(actualBase.c_str(), requestedBase.c_str()) != 0)
            {
                CloseHandle(proc);
                attempt.outcome = KillOutcome::IdentityMismatch;
                report.attempts.push_back(attempt);
                continue;
            }

            if (IsProtectedImage(actualBase))
            {
                CloseHandle(proc);
                attempt.outcome = KillOutcome::Protected;
                report.attempts.push_back(attempt);
                continue;
            }

            if (dryRun)
            {
                CloseHandle(proc);
                attempt.outcome = KillOutcome::WouldKill;
                report.attempts.push_back(attempt);
                continue;
            }

            const bool ok = TerminateProcessViaDriver(pid);

            CloseHandle(proc);

            if (ok)
            {
                attempt.outcome = KillOutcome::Killed;
                ++kills;
                report.killed = (int)kills;
            }
            else
            {
                attempt.outcome = KillOutcome::DriverFailed;
            }

            report.attempts.push_back(attempt);

            // Space kills out; a tight loop of terminate IOCTLs is both noisy
            // and a good way to catch a scheduler hiccup.
            Sleep(150 + (GetTickCount() % 200));
        }
    }
}