proxy-god/proxy_chain_manager/firewall.py

"""
Windows Firewall enforcement — kill-switch style.

When engaged:
  - Default outbound policy → BLOCK (all profiles)
  - Allow rules for: GOST, NordVPN stack, this app, loopback, DHCP
  - Everything else is denied outbound — apps that don't go through the
    local proxy simply can't reach the internet.

When disengaged:
  - Remove our rules, restore default outbound → ALLOW
"""
from __future__ import annotations

import ctypes
import logging
import subprocess
import sys
from pathlib import Path

from .paths import gost_exe_path
from .vpn_detect import expand_vpn_executables

log = logging.getLogger(__name__)

RULE_PREFIX = "PCM_"


def is_admin() -> bool:
    try:
        return ctypes.windll.shell32.IsUserAnAdmin() != 0  # type: ignore[attr-defined]
    except Exception:
        return False


def request_admin_relaunch() -> bool:
    """Re-launch the current process elevated. Returns True if launched, False on error."""
    if is_admin():
        return False
    try:
        exe = sys.executable
        args = " ".join(f'"{a}"' for a in sys.argv)
        ret = ctypes.windll.shell32.ShellExecuteW(  # type: ignore[attr-defined]
            None, "runas", exe, args, None, 1
        )
        return ret > 32
    except Exception:
        return False


def _run(args: list[str], check: bool = False) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        args, capture_output=True, text=True, timeout=30,
        creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
    )


def _add_rule(name: str, **kw: str) -> bool:
    full = RULE_PREFIX + name
    args = ["netsh", "advfirewall", "firewall", "add", "rule", f"name={full}"]
    for k, v in kw.items():
        args.append(f"{k}={v}")
    r = _run(args)
    ok = r.returncode == 0
    if not ok:
        log.warning("Firewall rule %s failed: %s", full, (r.stdout or "") + (r.stderr or ""))
    return ok


def _delete_rules() -> None:
    r = _run(["netsh", "advfirewall", "firewall", "show", "rule", "name=all", "dir=out"])
    lines = (r.stdout or "").splitlines()
    names: set[str] = set()
    for ln in lines:
        if ln.startswith("Rule Name:"):
            n = ln.split(":", 1)[1].strip()
            if n.startswith(RULE_PREFIX):
                names.add(n)
    for n in names:
        _run(["netsh", "advfirewall", "firewall", "delete", "rule", f"name={n}"])
    log.info("Deleted %d PCM firewall rules", len(names))


def _set_outbound_policy(action: str) -> None:
    """action = 'blockinbound,blockoutbound' or 'blockinbound,allowoutbound'."""
    for profile in ("domainprofile", "privateprofile", "publicprofile"):
        _run(["netsh", "advfirewall", "set", profile, "firewallpolicy", action])


def _resolve_self_exe() -> Path:
    if getattr(sys, "frozen", False):
        return Path(sys.executable).resolve()
    return Path(sys.executable).resolve()


def engage(gost_path: Path | None = None) -> tuple[bool, str]:
    """Activate kill-switch firewall. Returns (success, message)."""
    if not is_admin():
        return False, "Need admin privileges for firewall enforcement."

    gost = gost_path or gost_exe_path()
    self_exe = _resolve_self_exe()
    vpn_exes = expand_vpn_executables()

    _delete_rules()

    # Allow loopback (apps → GOST on 127.0.0.1)
    _add_rule("Loopback", dir="out", action="allow",
              remoteip="127.0.0.0/8", protocol="any")

    # Allow GOST outbound (it connects to the proxy hops)
    _add_rule("GOST", dir="out", action="allow",
              program=f'"{gost}"', protocol="any")

    # Allow this application outbound (fetches proxy lists, validates)
    _add_rule("Self", dir="out", action="allow",
              program=f'"{self_exe}"', protocol="any")

    # If running from python.exe, also allow pythonw.exe
    if self_exe.name.lower() in ("python.exe", "pythonw.exe"):
        for sibling in ("python.exe", "pythonw.exe"):
            p = self_exe.with_name(sibling)
            if p.is_file():
                _add_rule(f"Python_{sibling}", dir="out", action="allow",
                          program=f'"{p}"', protocol="any")

    # Allow VPN client executables (Nord, WireGuard, OpenVPN, etc.)
    for i, vpath in enumerate(vpn_exes):
        _add_rule(f"VPN_{i}", dir="out", action="allow",
                  program=f'"{vpath}"', protocol="any")

    # Allow DHCP (or you lose your adapter)
    _add_rule("DHCP", dir="out", action="allow",
              protocol="udp", remoteport="67,68")

    # Allow DNS (GOST needs to resolve proxy hostnames)
    _add_rule("DNS_UDP", dir="out", action="allow",
              protocol="udp", remoteport="53")
    _add_rule("DNS_TCP", dir="out", action="allow",
              protocol="tcp", remoteport="53")

    # Set default outbound to BLOCK
    _set_outbound_policy("blockinbound,blockoutbound")

    log.info("Firewall kill-switch engaged. %d VPN exes whitelisted.", len(vpn_exes))
    return True, f"Kill-switch ON. {len(vpn_exes)} VPN client(s) whitelisted."


def disengage() -> tuple[bool, str]:
    """Remove kill-switch, restore normal outbound."""
    if not is_admin():
        return False, "Need admin to remove firewall rules."
    _set_outbound_policy("blockinbound,allowoutbound")
    _delete_rules()
    log.info("Firewall kill-switch disengaged.")
    return True, "Kill-switch OFF. Normal outbound restored."


def is_engaged() -> bool:
    """Quick check: are our rules present?"""
    r = _run(["netsh", "advfirewall", "firewall", "show", "rule", f"name={RULE_PREFIX}GOST", "dir=out"])
    return RULE_PREFIX in (r.stdout or "")


def emergency_disengage() -> None:
    """Best-effort kill-switch removal for atexit/signal handlers.

    Unlike ``disengage()``, this never raises and does not require prior
    admin check — it simply tries, logs the outcome, and returns.  Safe to
    call from atexit or signal handlers where exceptions must not propagate.

    No-op when our rules are not present (avoids touching firewall policy on
    every process exit).
    """
    if not is_engaged():
        return
    try:
        _set_outbound_policy("blockinbound,allowoutbound")
        _delete_rules()
        log.info("emergency_disengage: firewall rules removed.")
    except Exception as exc:  # noqa: BLE001
        log.warning("emergency_disengage failed: %s", exc)