yt-downloader/yt_downloader/vpn.py

"""Isolated Mullvad WireGuard tunnel for the yt-dlp download phase.

The download phase is the only place yt-downloader talks to YouTube. When a
WireGuard (Mullvad) config is configured, ``yt-downloader download`` re-runs
itself inside a dedicated network namespace whose only egress is the WireGuard
tunnel, so YouTube traffic never leaves through the normal network.

Layout inside the namespace:

    +-----------+   veth   +--------------------------------------+
    | host      | eth0     | netns "mullvad"                      |
    | veth-mv0  | 10.66..2 |   mv0 (WireGuard)  <- default route  |
    | 10.66..1  +----------+   DNS -> Mullvad (through tunnel)    |
    +-----------+          +--------------------------------------+

The namespace has no fallback egress: if the tunnel is down the default route
(via mv0) is dead, so traffic is dropped rather than leaked onto the host's
network.

Privileged steps are delegated to ``sudo``. sudo caches its credential, so an
interactive run prompts once; unattended runs need the passwordless sudoers
entries described in the README.
"""

from __future__ import annotations

import configparser
import getpass
import logging
import re
import shutil
import subprocess
import sys
from pathlib import Path

log = logging.getLogger(__name__)

# host <-> namespace point-to-point veth link (a /30)
_LINK_GUEST = "eth0"
_HOST_ADDR = "10.66.66.1/30"
_GUEST_ADDR = "10.66.66.2/30"
_GUEST_NET = "10.66.66.0/30"
_GUEST_GW = "10.66.66.1"
FALLBACK_DNS = "10.64.0.1"  # Mullvad DNS

# env var set on the inner re-run so it doesn't wrap itself again
INNER_ENV = "YTDL_VPN"


# -- configuration helpers ------------------------------------------------

def configured(cfg) -> bool:
    return bool(getattr(cfg, "vpn_wireguard", None))


def netns_name(cfg) -> str:
    return getattr(cfg, "vpn_namespace", None) or "mullvad"


def iface_name(cfg) -> str:
    return getattr(cfg, "vpn_interface", None) or "mv0"


def conf_path(cfg) -> Path:
    wg = getattr(cfg, "vpn_wireguard", None)
    if not wg:
        raise RuntimeError(
            "no WireGuard config configured: set [vpn] wireguard in the config "
            "(run 'yt-downloader init' and edit the file)"
        )
    return Path(wg).expanduser()


def _host_link(iface: str) -> str:
    return f"veth-{iface}"[:15]


def parse_conf(cfg) -> dict:
    """Read [Interface] / [Peer] from the WireGuard config."""
    path = conf_path(cfg)
    if not path.exists():
        raise FileNotFoundError(
            f"WireGuard config not found: {path} "
            "(set [vpn] wireguard or place the Mullvad .conf there)"
        )
    parser = configparser.ConfigParser(interpolation=None)
    try:
        parser.read(path)
    except configparser.Error as exc:
        raise ValueError(f"{path}: cannot parse WireGuard config: {exc}") from exc
    if not parser.has_section("Interface") or not parser.has_section("Peer"):
        raise ValueError(f"{path}: expected [Interface] and [Peer] sections")
    interface = dict(parser.items("Interface"))
    peer = dict(parser.items("Peer"))
    dns = [
        part.strip()
        for part in re.split(r"[\s,]+", interface.get("dns", "").strip())
        if part.strip()
    ]
    return {
        "interface": interface,
        "peer": peer,
        "dns": dns,
        "endpoint": peer.get("endpoint", "").strip(),
    }


def _strip_dns(path: Path) -> str:
    lines = [ln for ln in path.read_text(encoding="utf-8").splitlines()
             if not re.match(r"\s*DNS\s*=", ln)]
    return "\n".join(lines) + "\n"


# -- low-level sudo plumbing ---------------------------------------------

def _run(cmd: list[str], check: bool = True, input: bytes | None = None) -> subprocess.CompletedProcess:
    proc = subprocess.run(cmd, input=input, check=False)
    if check and proc.returncode != 0:
        raise RuntimeError(f"command failed (exit {proc.returncode}): {' '.join(cmd)}")
    return proc


def _sudo(args: list[str], check: bool = True, input: bytes | None = None) -> subprocess.CompletedProcess:
    return _run(["sudo", *args], check=check, input=input)


def _sudo_out(args: list[str]) -> tuple[int, str]:
    proc = subprocess.run(["sudo", *args], capture_output=True, text=True)
    return proc.returncode, (proc.stdout + proc.stderr).strip()


# -- lifecycle -------------------------------------------------------------

def is_up(cfg) -> bool:
    ns, iface = netns_name(cfg), iface_name(cfg)
    rc, out = _sudo_out(["ip", "netns", "exec", ns, "wg", "show", iface])
    return rc == 0 and "no such device" not in out.lower()


def handshake_age(cfg) -> int | None:
    """Seconds since the last WireGuard handshake, or None if none/unknown."""
    ns, iface = netns_name(cfg), iface_name(cfg)
    rc, out = _sudo_out(["ip", "netns", "exec", ns, "wg", "show", iface])
    if rc != 0:
        return None
    for line in out.splitlines():
        m = re.search(r"latest handshake:\s*(.+)", line)
        if not m:
            continue
        text = m.group(1).strip().lower()
        if text == "no handshake":
            return None
        m2 = re.match(r"(\d+)\s+seconds? ago", text)
        if m2:
            return int(m2.group(1))
    return None


def has_handshake(cfg, max_age: int = 180) -> bool:
    """True when the tunnel has a recent handshake (i.e. is actually usable)."""
    age = handshake_age(cfg)
    return age is not None and age <= max_age


def up(cfg) -> None:
    ns, iface = netns_name(cfg), iface_name(cfg)
    host_link = _host_link(iface)
    for tool in ("ip", "wg", "wg-quick"):
        if not shutil.which(tool):
            raise RuntimeError(
                f"required tool not found: {tool!r} (install wireguard-tools / iproute2)"
            )
    info = parse_conf(cfg)
    dns_servers = info["dns"] or [FALLBACK_DNS]

    # 1. namespace DNS — `ip netns exec` binds this over /etc/resolv.conf
    _sudo(["mkdir", "-p", f"/etc/netns/{ns}"])
    resolv = "".join(f"nameserver {d}\n" for d in dns_servers)
    _sudo(["sh", "-c", f"cat > /etc/netns/{ns}/resolv.conf"], input=resolv.encode())

    # 2. namespace + host <-> namespace veth link
    _sudo(["ip", "netns", "add", ns], check=False)
    _sudo(["ip", "link", "add", host_link, "type", "veth",
           "peer", "name", _LINK_GUEST, "netns", ns], check=False)
    _sudo(["ip", "link", "set", host_link, "up"])
    _sudo(["ip", "addr", "add", _HOST_ADDR, "dev", host_link], check=False)
    _sudo(["ip", "netns", "exec", ns, "ip", "link", "set", "lo", "up"])
    _sudo(["ip", "netns", "exec", ns, "ip", "addr", "add", _GUEST_ADDR, "dev", _LINK_GUEST], check=False)
    _sudo(["ip", "netns", "exec", ns, "ip", "link", "set", _LINK_GUEST, "up"])
    # pre-tunnel default route so the WireGuard endpoint is reachable while the
    # tunnel is still down; wg-quick replaces it with the tunnel default below
    _sudo(["ip", "netns", "exec", ns, "ip", "route", "add", "default", "via", _GUEST_GW], check=False)

    # 3. host forwards + NATs the namespace so it can reach the wg endpoint
    _sudo(["sysctl", "-w", "net.ipv4.ip_forward=1"])
    _sudo(["iptables", "-t", "nat", "-C", "POSTROUTING",
           "-s", _GUEST_NET, "-j", "MASQUERADE"], check=False)
    _sudo(["iptables", "-t", "nat", "-A", "POSTROUTING",
           "-s", _GUEST_NET, "-j", "MASQUERADE"])

    # 4. normalized wg config (DNS handled above, not by wg-quick) + tunnel up
    cleaned = _strip_dns(conf_path(cfg))
    _sudo(["sh", "-c", f"cat > /etc/wireguard/{iface}.conf"], input=cleaned.encode())
    _sudo(["chmod", "600", f"/etc/wireguard/{iface}.conf"])
    _sudo(["ip", "netns", "exec", ns, "wg-quick", "up", iface])

    log.info("Mullvad tunnel up: namespace=%s interface=%s dns=%s",
             ns, iface, ",".join(dns_servers))


def down(cfg) -> None:
    ns, iface = netns_name(cfg), iface_name(cfg)
    host_link = _host_link(iface)
    _sudo(["ip", "netns", "exec", ns, "wg-quick", "down", iface], check=False)
    _sudo(["ip", "link", "del", host_link], check=False)
    _sudo(["ip", "netns", "del", ns], check=False)
    _sudo(["iptables", "-t", "nat", "-D", "POSTROUTING",
           "-s", _GUEST_NET, "-j", "MASQUERADE"], check=False)
    _sudo(["rm", "-f", f"/etc/wireguard/{iface}.conf"])
    _sudo(["rm", "-rf", f"/etc/netns/{ns}"])


def _poke(cfg) -> None:
    """Send a tiny packet through the tunnel to force a WireGuard handshake."""
    ns = netns_name(cfg)
    code = (
        "import socket;"
        "s=socket.socket(socket.AF_INET,socket.SOCK_DGRAM);"
        "s.sendto(b'\\x00'*32, ('10.64.0.1', 53))"
    )
    _sudo(["ip", "netns", "exec", ns, "python3", "-c", code], check=False)


def wait_for_handshake(cfg, timeout: float = 20.0) -> bool:
    """Wait for a handshake, actively routing packets to trigger one.

    WireGuard only initiates a handshake once traffic flows through the
    tunnel, so an idle interface would never handshake on its own.
    """
    import time
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        _poke(cfg)  # force traffic -> handshake
        if has_handshake(cfg):
            return True
        time.sleep(0.5)
    return has_handshake(cfg)


def _diagnose(cfg) -> str:
    """Collect tunnel state to explain why a handshake isn't happening."""
    ns, iface = netns_name(cfg), iface_name(cfg)
    parts = []
    for title, args in (
        ("wg", ["ip", "netns", "exec", ns, "wg", "show", iface]),
        ("routes", ["ip", "netns", "exec", ns, "ip", "route", "show", "default"]),
        ("rules", ["ip", "netns", "exec", ns, "ip", "rule", "show"]),
    ):
        rc, out = _sudo_out(args)
        parts.append(f"--- {title} ---\n{out or '(none)'}")
    return "\n".join(parts)


def ensure_up(cfg) -> None:
    if has_handshake(cfg):
        return
    log.info("Mullvad tunnel is down or not handshaken — bringing it up")
    down(cfg)  # converge from any stale half-configured state
    up(cfg)
    if not wait_for_handshake(cfg):
        log.error("no WireGuard handshake within timeout — tunnel state:\n%s",
                  _diagnose(cfg))
        down(cfg)
        raise RuntimeError(
            "WireGuard tunnel came up but no handshake was established — "
            "see tunnel state above. If 'transfer' shows sent bytes but 0 "
            "received, the endpoint is unreachable (try a different Mullvad "
            "server); if 0 sent, the namespace has no route out. "
            "Refusing to download over a dead tunnel."
        )


def status_text(cfg) -> str:
    ns, iface = netns_name(cfg), iface_name(cfg)
    rc, out = _sudo_out(["ip", "netns", "exec", ns, "wg", "show", iface])
    if rc != 0:
        return "tunnel: DOWN (namespace or WireGuard interface missing)"
    lines = ["tunnel: UP", f"namespace : {ns}"]
    lines.extend(line for line in out.splitlines())
    age = handshake_age(cfg)
    lines.append("handshake : " + ("usable" if has_handshake(cfg)
                                   else ("none" if age is None else f"{age}s ago (stale)")))
    _, route = _sudo_out(["ip", "netns", "exec", ns, "ip", "route", "show", "default"])
    lines.append("default   : " + (route or "(none)"))
    _, dns = _sudo_out(["cat", f"/etc/netns/{ns}/resolv.conf"])
    lines.append("dns       : " + (dns.replace("\n", " ").strip() or "(none)"))
    return "\n".join(lines)


# -- running downloads inside the tunnel ----------------------------------

def run_download_in_ns(cfg, argv: list[str]) -> int:
    """Re-run ``yt-downloader download <argv>`` inside the tunnel namespace,
    dropping back to the invoking user so downloaded files stay theirs."""
    ensure_up(cfg)
    ns = netns_name(cfg)
    user = getpass.getuser()
    cmd = [sys.executable, "-m", "yt_downloader", *argv]
    wrapped = ["sudo", "ip", "netns", "exec", ns,
               "sudo", "-u", user, "env", f"{INNER_ENV}=1", *cmd]
    log.info("running download inside Mullvad namespace %r", ns)
    return _run(wrapped, check=False).returncode