main.cpp

#include <windows.h>
#include <tlhelp32.h>
#include <thread>
#include <atomic>
#include <mutex>
#include <string>
#include <chrono>
#include "windivert.h"

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

// ── Глобальные переменные ──────────────────────────
HWND              g_hwnd    = nullptr;
HWND              g_status  = nullptr;
HWND              g_btn     = nullptr;

HANDLE            g_divert  = INVALID_HANDLE_VALUE;
std::atomic<bool> g_running { false };
std::atomic<bool> g_alive   { true  }; // для WatcherLoop
std::mutex        g_mutex;             // защита от гонки
std::thread       g_worker;
std::thread       g_watcher;

// ── Статус — безопасно через копию строки ──────────
void SetStatus(const std::wstring& text) {
    // Копируем строку в кучу — живёт до обработки сообщения
    std::wstring* copy = new std::wstring(text);
    PostMessage(g_hwnd, WM_USER + 1, 0, (LPARAM)copy);
}

// ── Проверка процессов ─────────────────────────────
bool IsProcessRunning(const wchar_t* name) {
    HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (snap == INVALID_HANDLE_VALUE) return false;

    PROCESSENTRY32W pe{ sizeof(pe) };
    bool found = false;

    if (Process32FirstW(snap, &pe))
        do {
            if (_wcsicmp(pe.szExeFile, name) == 0)
                { found = true; break; }
        } while (Process32NextW(snap, &pe));

    CloseHandle(snap); // всегда закрываем handle
    return found;
}

bool IsFirefoxPID(DWORD pid) {
    if (pid == 0) return false; // защита от нулевого PID

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

    PROCESSENTRY32W pe{ sizeof(pe) };
    bool found = false;

    if (Process32FirstW(snap, &pe))
        do {
            if (pe.th32ProcessID == pid &&
                _wcsicmp(pe.szExeFile, L"firefox.exe") == 0)
                { found = true; break; }
        } while (Process32NextW(snap, &pe));

    CloseHandle(snap);
    return found;
}

// ── Цикл перехвата пакетов ─────────────────────────
void DivertLoop() {
    // Буфер на стеке — не нужен delete[]
    BYTE buf[65535];
    UINT len = 0;
    WINDIVERT_ADDRESS addr{};

    while (g_running) {
        if (!WinDivertRecv(g_divert, buf, sizeof(buf), &len, &addr))
            continue;

        // Пропускаем не-Firefox трафик без изменений
        if (!IsFirefoxPID(addr.Flow.ProcessId)) {
            WinDivertSend(g_divert, buf, len, nullptr, &addr);
            continue;
        }

        PWINDIVERT_IPHDR  ip      = nullptr;
        PWINDIVERT_TCPHDR tcp     = nullptr;
        PVOID             payload = nullptr;
        UINT              payLen  = 0;

        WinDivertHelperParsePacket(
            buf, len,
            &ip, nullptr, nullptr, nullptr,
            nullptr, &tcp, nullptr,
            &payload, &payLen
        );

        // ip и tcp — не наша память, не удаляем — просто указатели внутрь buf
        if (ip && tcp && payLen > 0) {
            BYTE realTTL = ip->TTL;

            // Обманный пакет — умрёт до DPI провайдера
            ip->TTL = 5;
            WinDivertHelperCalcChecksums(buf, len, &addr, 0);
            WinDivertSend(g_divert, buf, len, nullptr, &addr);

            // Восстанавливаем — реальный пакет
            ip->TTL = realTTL;
        }

        WinDivertHelperCalcChecksums(buf, len, &addr, 0);
        WinDivertSend(g_divert, buf, len, nullptr, &addr);
    }
}

// ── Запуск / Остановка движка ─────────────────────
void StartEngine() {
    std::lock_guard<std::mutex> lock(g_mutex); // блокируем гонку

    if (g_running) return;

    g_divert = WinDivertOpen(
        "outbound and tcp and "
        "(tcp.DstPort == 443 or tcp.DstPort == 80) and "
        "tcp.PayloadLength > 0",
        WINDIVERT_LAYER_NETWORK, 0, 0
    );

    if (g_divert == INVALID_HANDLE_VALUE) {
        SetStatus(L"❌ Ошибка — запусти от администратора");
        return;
    }

    g_running = true;
    g_worker  = std::thread(DivertLoop);

    SetStatus(L"✅ Работает");
    SetWindowTextW(g_btn, L"Остановить");
}

void StopEngine() {
    std::lock_guard<std::mutex> lock(g_mutex); // блокируем гонку

    if (!g_running) return;

    g_running = false;

    // Закрываем handle — это разблокирует WinDivertRecv в потоке
    if (g_divert != INVALID_HANDLE_VALUE) {
        WinDivertClose(g_divert);
        g_divert = INVALID_HANDLE_VALUE;
    }

    // Ждём завершения потока — нет утечки потока
    if (g_worker.joinable())
        g_worker.join();

    SetStatus(L"⏸ Остановлен");
    SetWindowTextW(g_btn, L"Включить вручную");
}

// ── Слежка за Firefox — с корректным выходом ───────
void WatcherLoop() {
    bool was = false;

    while (g_alive) { // выходим при закрытии программы
        bool is = IsProcessRunning(L"firefox.exe");

        if  (is && !was) StartEngine();
        if (!is &&  was) StopEngine();

        was = is;
        std::this_thread::sleep_for(std::chrono::milliseconds(1000));
    }
}

// ── Обработчик окна ────────────────────────────────
LRESULT CALLBACK WndProc(HWND hwnd, UINT msg,
                          WPARAM wp, LPARAM lp)
{
    switch (msg) {
    case WM_CREATE:
        g_status = CreateWindowW(
            L"STATIC", L"⏳ Ожидание Firefox...",
            WS_CHILD | WS_VISIBLE | SS_CENTER,
            10, 20, 360, 30,
            hwnd, nullptr, GetModuleHandle(nullptr), nullptr);

        g_btn = CreateWindowW(
            L"BUTTON", L"Включить вручную",
            WS_CHILD | WS_VISIBLE | BS_PUSHBUTTON,
            120, 65, 155, 35,
            hwnd, (HMENU)1,
            GetModuleHandle(nullptr), nullptr);

        g_watcher = std::thread(WatcherLoop);
        break;

    case WM_COMMAND:
        if (LOWORD(wp) == 1)
            g_running ? StopEngine() : StartEngine();
        break;

    case WM_USER + 1: {
        // Получаем строку из кучи и сразу удаляем — нет утечки
        std::wstring* text = reinterpret_cast<std::wstring*>(lp);
        SetWindowTextW(g_status, text->c_str());
        delete text; // освобождаем память
        break;
    }

    case WM_DESTROY:
        g_alive = false; // останавливаем WatcherLoop
        StopEngine();
        if (g_watcher.joinable())
            g_watcher.join(); // ждём завершения
        PostQuitMessage(0);
        break;
    }

    return DefWindowProcW(hwnd, msg, wp, lp);
}

// ── Точка входа ────────────────────────────────────
int WINAPI WinMain(HINSTANCE hInst, HINSTANCE,
                   LPSTR, int nShow)
{
    WNDCLASSW wc{};
    wc.lpfnWndProc   = WndProc;
    wc.hInstance     = hInst;
    wc.lpszClassName = L"DPIApp";
    wc.hbrBackground = (HBRUSH)(COLOR_WINDOW + 1);
    wc.hCursor       = LoadCursor(nullptr, IDC_ARROW);
    RegisterClassW(&wc);

    g_hwnd = CreateWindowW(
        L"DPIApp", L"DPI Bypass",
        WS_OVERLAPPEDWINDOW & ~WS_THICKFRAME & ~WS_MAXIMIZEBOX,
        CW_USEDEFAULT, CW_USEDEFAULT, 400, 140,
        nullptr, nullptr, hInst, nullptr);

    ShowWindow(g_hwnd, nShow);
    UpdateWindow(g_hwnd);

    MSG msg{};
    while (GetMessageW(&msg, nullptr, 0, 0)) {
        TranslateMessage(&msg);
        DispatchMessageW(&msg);
    }
    return (int)msg.wParam;
}