HexUnlocked/src/unlocker/main.cpp

#include "win-platform.h"
#include "tray-icon.h"
#include "utils.h"
#include "spoofer.h"
#include "log-sink-cout.h"
#include "log-sink-file.h"
#include "proxy-configurator.h"
#include "cache-cleaner.h"

#include <seallib/log.h>
#include <tinymitm/proxy.h>

#include <memory>

#include <processenv.h>
#include <cassert>
#include <processthreadsapi.h>
#include <handleapi.h>

bool running = true;
TinyMITMProxy* proxy = nullptr;
ProxyConfigurator* conf = nullptr;
CacheCleaner* cleaner = nullptr;

BOOL WINAPI consoleHandler(DWORD dwType)
{
    if (dwType == CTRL_C_EVENT || dwType == CTRL_CLOSE_EVENT || dwType == CTRL_LOGOFF_EVENT ||
        dwType == CTRL_SHUTDOWN_EVENT)
    {
        if (conf) conf->clearProxy();
        running = false;
        return TRUE;
    }
    return FALSE;
}

void StartWatchdog()
{
    DWORD pid = GetCurrentProcessId();

    /*
        this probably looks so sketchy for some antiviruses
        should test, probably
    */
    std::string psCommand =
        std::format("powershell.exe -WindowStyle Hidden -Command \""
                    "$proc = Get-Process -Id {} -ErrorAction SilentlyContinue;"
                    "if ($proc) {{ $proc.WaitForExit(); }}"
                    "$reg = 'HKCU:\\Software\\Microsoft\\Windows\\CurrentVersion\\Internet Settings';"
                    "Set-ItemProperty -Path $reg -Name ProxyEnable -Value 0;"
                    "Set-ItemProperty -Path $reg -Name ProxyServer -Value '';"
                    "rundll32.exe user32.dll,UpdatePerUserSystemParameters;\"",
                    pid);

    STARTUPINFOA si;
    PROCESS_INFORMATION pi;
    RtlZeroMemory(&si, sizeof(si));
    RtlZeroMemory(&pi, sizeof(pi));
    si.cb = sizeof(si);

    CreateProcessA(NULL, const_cast<char*>(psCommand.c_str()), NULL, NULL, FALSE, CREATE_NO_WINDOW, NULL, NULL, &si,
                   &pi);

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

bool run()
{
    /*
        mutex check before running
    */
    HANDLE mtx = CreateMutexA(NULL, TRUE, "Global\\HexUnlocked");
    if (mtx == NULL) return false;
    if (GetLastError() == ERROR_ALREADY_EXISTS)
    {
        MessageBoxA(NULL, "Program is already running", "Hex: Unlocked", MB_OK | MB_ICONERROR | MB_ICONINFORMATION);
        CloseHandle(mtx);
        return false;
    }

    /*
        init
    */
    seallib::Logger mainLog("Main");
    mainLog.addSink(std::make_shared<ConOutSink>());
    mainLog.addSink(FileSink::getSharedInstance());

    mainLog.info("Init");

    mainLog.info("Setting up cache cleaner");
    cleaner = new CacheCleaner();
    cleaner->init();

    mainLog.info("Creating proxy");
    TinyMITMConfig mitmConfig;
    mitmConfig.port = 44444;
    mitmConfig.threadCount = 255;
    mitmConfig.caName = utils::randomizeString(32);
    proxy = new TinyMITMProxy(mitmConfig);

    mainLog.info("Creating configurator and setting proxy addr");
    conf = new ProxyConfigurator();
    conf->setProxy("127.0.0.1", 44444);

    mainLog.info("Instantiating spoofer and registering proxy listeners");
    Spoofer* spoofer = new Spoofer();
    spoofer->registerListeners(proxy);

    mainLog.info("Starting proxy");
    if (!proxy->init())
    {
        mainLog.error("Failed to initialize proxy, shutting down");
        proxy->shutdown();
        cleaner->shutdown();
        return false;
    }

    mainLog.info("Setting up system tray icon");

    TrayIcon tray;
    tray.setExitCallback([&]() {
        running = false;
        mainLog.info("Exit requested from system tray.");
    });

    if (!tray.init())
    {
        mainLog.error("Failed to initialize system tray icon.");
        return false;
    }

    HWND consoleWnd = GetConsoleWindow();
    if (consoleWnd)
    {
        HMENU menuHandle = GetSystemMenu(consoleWnd, FALSE);
        if (menuHandle) RemoveMenu(menuHandle, SC_CLOSE, MF_BYCOMMAND);
    }

    mainLog.info("Proxy running, Ctrl+C to stop. Check system tray for options.");
    mainLog.info("Go to https://dbd.neru.rip/ for settings.");
    while (running)
    {
        tray.processMessages();
        std::this_thread::sleep_for(std::chrono::milliseconds(50));
    }

    mainLog.info("Shutting down...");
    tray.shutdown();
    proxy->shutdown();
    cleaner->shutdown();

    delete proxy;
    proxy = nullptr;
    delete conf;
    conf = nullptr;

    ReleaseMutex(mtx);
    CloseHandle(mtx);

    return true;
}

int WINAPI WinMain(HINSTANCE /*hInstance*/, HINSTANCE /*hPrevInstance*/, LPSTR /*lpCmdLine*/, int /*nShowCmd*/)
{
    AllocConsole();
    SetConsoleTitleA(utils::randomizeString(32).c_str());

    /*
        ansi seequences
    */
    HANDLE outHandle = GetStdHandle(STD_OUTPUT_HANDLE);
    if (outHandle == INVALID_HANDLE_VALUE) return 1;

    DWORD conMode = 0;
    if (!GetConsoleMode(outHandle, &conMode)) return 1;

    conMode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING;
    SetConsoleMode(outHandle, conMode);

    /*
        io
    */
    FILE* dummy;
    freopen_s(&dummy, "CONIN$", "r", stdin);
    freopen_s(&dummy, "CONOUT$", "w", stdout);
    freopen_s(&dummy, "CONOUT$", "w", stderr);

    /*
        handlers for cleaning proxy conf on exit / crash / taskkill / whatever
    */
    SetConsoleCtrlHandler(consoleHandler, TRUE);
    StartWatchdog();

    // __try
    // {
    if (!run()) return 1;
    // }
    // __except (EXCEPTION_EXECUTE_HANDLER)
    // {
    //     if (conf) conf->clearProxy();
    // }

    return 0;
}