proxy-god/proxy_chain_manager/gost_util.py

from __future__ import annotations

import hashlib
import io
import logging
import shutil
import subprocess
import zipfile
from pathlib import Path
from typing import Any

import httpx

from .paths import app_data_dir, gost_exe_path

log = logging.getLogger(__name__)

GOST_RELEASE_ZIP = (
    "https://github.com/go-gost/gost/releases/download/v3.2.6/"
    "gost_3.2.6_windows_amd64.zip"
)

# Pinned SHA256 of the release zip.  Update this constant when bumping GOST_RELEASE_ZIP.
# Verified 2026-05-21 against https://github.com/go-gost/gost/releases/download/v3.2.6/
GOST_RELEASE_ZIP_SHA256 = "32f4edf3d94b622e67f1979f6f5de82dac62abc0977772cf96215dd199ef7e7b"



def _add_defender_exclusion(path: Path) -> None:
    """Add Windows Defender exclusion so GOST is not quarantined."""
    try:
        # ExclusionPath covers gost.exe under this folder; avoid invalid combined params on older builds.
        folder = str(path.parent).replace("'", "''")
        subprocess.run(
            [
                "powershell", "-NoProfile", "-NonInteractive", "-Command",
                f"Add-MpPreference -ExclusionPath '{folder}' -ErrorAction SilentlyContinue",
            ],
            capture_output=True,
            timeout=30,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
        )
        log.info("Defender exclusion added for %s", path.parent)
    except Exception as exc:
        log.warning("Defender exclusion failed (non-fatal): %s", exc)


def _verify_zip_sha256(data: bytes, expected: str) -> None:
    """Raise RuntimeError if SHA256 of *data* does not match *expected* (hex)."""
    actual = hashlib.sha256(data).hexdigest().lower()
    if actual != expected.lower():
        raise RuntimeError(
            f"GOST zip SHA256 mismatch — possible supply-chain attack!\n"
            f"  expected: {expected.lower()}\n"
            f"  actual:   {actual}\n"
            "Delete the downloaded zip and retry, or update GOST_RELEASE_ZIP_SHA256."
        )
    log.info("GOST zip SHA256 OK: %s", actual)


def ensure_gost(target: Path | None = None) -> Path:
    exe = target or gost_exe_path()
    if exe.is_file() and exe.stat().st_size > 10_000:
        return exe
    exe.parent.mkdir(parents=True, exist_ok=True)
    # Add exclusion BEFORE downloading so Defender doesn't nuke it on write
    _add_defender_exclusion(exe)
    log.info("Downloading GOST %s", GOST_RELEASE_ZIP)
    with httpx.Client(timeout=120.0, follow_redirects=True) as c:
        r = c.get(GOST_RELEASE_ZIP)
        r.raise_for_status()
        data = r.content
    if len(data) < 64 or data[:2] != b"PK":
        raise RuntimeError("Downloaded GOST zip looks invalid (not a zip).")
    # Integrity check before extraction
    _verify_zip_sha256(data, GOST_RELEASE_ZIP_SHA256)
    with zipfile.ZipFile(io.BytesIO(data), "r") as z:
        names = [n for n in z.namelist() if n.lower().endswith("gost.exe")]
        if not names:
            raise RuntimeError("gost.exe not found in release zip")
        with z.open(names[0]) as src, open(exe, "wb") as dst:
            shutil.copyfileobj(src, dst)
    # Add exclusion again after write in case Defender scanned during extraction
    _add_defender_exclusion(exe)
    log.info("GOST installed at %s", exe)
    return exe


def build_gost_cmd(gost: Path, listen_http: str, forwards: list[str]) -> list[str]:
    args = [str(gost), "-L", f"http://{listen_http}"]
    for f in forwards:
        args.extend(["-F", f])
    return args


def gost_log_path() -> Path:
    return app_data_dir() / "gost.log"


def popen_no_window(args: list[str]) -> subprocess.Popen:
    # stderr → rolling log file (not a pipe — pipes deadlock if unread; files never block).
    # stdout is also captured to the same file since GOST v3 mixes output channels.
    cr = getattr(subprocess, "CREATE_NO_WINDOW", 0)
    log_path = gost_log_path()
    try:
        _rotate_gost_log(log_path)
        stderr_dest: Any = open(log_path, "ab")  # noqa: WPS515 — intentional long-lived file handle
    except OSError:
        stderr_dest = subprocess.DEVNULL
    return subprocess.Popen(
        args,
        stdin=subprocess.DEVNULL,
        stdout=stderr_dest,
        stderr=stderr_dest,
        creationflags=cr,
    )


def _rotate_gost_log(path: Path, max_bytes: int = 512 * 1024) -> None:
    """Keep gost.log under max_bytes by truncating the oldest half when it grows too large."""
    try:
        if path.is_file() and path.stat().st_size > max_bytes:
            data = path.read_bytes()
            path.write_bytes(data[len(data) // 2:])
    except OSError:
        pass


def read_gost_log_tail(lines: int = 40) -> str:
    """Return the last N lines of gost.log for display in the UI."""
    try:
        text = gost_log_path().read_text(encoding="utf-8", errors="replace")
        return "\n".join(text.splitlines()[-lines:])
    except OSError:
        return "(gost.log not found)"


def terminate_process(proc: subprocess.Popen | None) -> None:
    if proc is None:
        return
    if proc.poll() is not None:
        return
    try:
        proc.terminate()
        proc.wait(timeout=5)
    except Exception as exc:
        log.debug("GOST terminate failed (%s) — escalating to kill.", exc)
        try:
            proc.kill()
        except Exception as kill_exc:
            log.warning("GOST kill also failed: %s", kill_exc)