awg-dash/tests/netns/netns_common.py

#!/usr/bin/env python3
"""netns_common.py — утилиты проверок netns-стенда (задача #21). Только stdlib.

Используется тестами (tests/netns/test_netns_stand.py) и run_tests.py.
Значения по умолчанию совпадают с tests/netns/lib.sh — переменные окружения
позволяют переопределить топологию и передаются run.sh в дочерние процессы.
"""

from __future__ import annotations

import json
import os
import pathlib
import subprocess
import time
from typing import Any, Iterable

# --- топология ---------------------------------------------------------------
CLIENT_NS = os.environ.get("AWG_NS_CLIENT", "awgt21-client")
GW_NS = os.environ.get("AWG_NS_GW", "awgt21-gw")
UPLINK_NS = os.environ.get("AWG_NS_UPLINK", "awgt21-uplink")

CLIENT_IP = os.environ.get("AWG_CLIENT_IP", "10.213.0.2")
GW_LAN_IP = os.environ.get("AWG_GW_LAN_IP", "10.213.0.1")
GW_WAN_IP = os.environ.get("AWG_GW_WAN_IP", "10.214.0.1")
UPLINK_IP = os.environ.get("AWG_UPLINK_IP", "10.214.0.2")

MOCK_PORT = int(os.environ.get("AWG_MOCK_PORT", "51820"))
LEAK_IP = os.environ.get("AWG_LEAK_IP", "198.51.100.9")
LEAK_PORT = int(os.environ.get("AWG_LEAK_PORT", "9999"))

# имена veth-интерфейсов (совпадают с tests/netns/lib.sh)
VETH_CLIENT = os.environ.get("AWG_VETH_CLIENT", "veth-c-gw")
VETH_GW_LAN = os.environ.get("AWG_VETH_GW_LAN", "veth-gw-lan")
VETH_GW_WAN = os.environ.get("AWG_VETH_GW_WAN", "veth-gw-wan")
VETH_UPLINK = os.environ.get("AWG_VETH_UPLINK", "veth-uplink")

FWMARK = int(os.environ.get("AWG_FWMARK", "51820"))
ROUTE_TABLE = int(os.environ.get("AWG_ROUTE_TABLE", "51820"))

WORKDIR = pathlib.Path(os.environ.get("AWG_NETNS_WORKDIR", "/tmp/awg-netns"))
MOCK_LOG = pathlib.Path(os.environ.get("AWG_MOCK_LOG", str(WORKDIR / "mock-endpoint.jsonl")))
MOCK_READY = pathlib.Path(os.environ.get("AWG_MOCK_READY", str(WORKDIR / "mock-endpoint.ready")))
MOCK_MODE_FILE = pathlib.Path(os.environ.get("AWG_MOCK_MODE_FILE", str(WORKDIR / "mock-mode")))
LEAK_LOG = pathlib.Path(os.environ.get("AWG_LEAK_LOG", str(WORKDIR / "leak-sink.jsonl")))
LEAK_READY = pathlib.Path(os.environ.get("AWG_LEAK_READY", str(WORKDIR / "leak-sink.ready")))

HERE = pathlib.Path(__file__).resolve().parent
PROBE = HERE / "probe.py"

NFT_TABLE = "awgtest"
C_MARK_OUT = "awgtest-mark-out"
C_MARK_FWD = "awgtest-mark-fwd"
C_DROP = "awgtest-killswitch-drop"
IPT_CHAIN_MARK = "awgtest-mark"
IPT_CHAIN_FWD = "awgtest-mark-fwd"
IPT_CHAIN_KS = "awgtest-ks"


class StandError(RuntimeError):
    """Ошибка стенда (не теста): топология/утилиты не отвечают как ожидается."""


def sh(args: Iterable[str], check: bool = True, timeout: float = 20.0) -> subprocess.CompletedProcess:
    """Запуск команды в корневом netns."""
    cp = subprocess.run(
        list(args),
        capture_output=True,
        text=True,
        timeout=timeout,
        check=False,
    )
    if check and cp.returncode != 0:
        raise StandError(f"команда {list(args)!r} завершилась с кодом {cp.returncode}: {cp.stderr.strip()}")
    return cp


def ns_sh(ns: str, args: Iterable[str], check: bool = True, timeout: float = 20.0) -> subprocess.CompletedProcess:
    """Запуск команды внутри указанного netns."""
    return sh(["ip", "netns", "exec", ns, *args], check=check, timeout=timeout)


def ns_exists(ns: str) -> bool:
    cp = sh(["ip", "netns", "list"], check=False)
    return any(line.split()[0] == ns for line in cp.stdout.splitlines() if line.strip())


def addrs(ns: str) -> str:
    return sh(["ip", "-n", ns, "-br", "addr", "show"]).stdout


def links(ns: str) -> str:
    return sh(["ip", "-n", ns, "-br", "link", "show"]).stdout


def link_is_up(ns: str, ifname: str) -> bool:
    cp = sh(["ip", "-n", ns, "-br", "link", "show", ifname], check=False)
    return cp.returncode == 0 and "UP" in cp.stdout


# --- netfilter: backend, правила, счётчики -----------------------------------
def fw_kind() -> str:
    """Backend netfilter в netns шлюза: nft|iptables (AWG_FW_BACKEND форсирует выбор)."""
    want = os.environ.get("AWG_FW_BACKEND", "auto")
    if want in ("nft", "iptables"):
        if ns_sh(GW_NS, ["sh", "-c", f"command -v {want}"], check=False).returncode == 0:
            return want
        return "none"
    if ns_sh(GW_NS, ["sh", "-c", "command -v nft"], check=False).returncode == 0:
        return "nft"
    if ns_sh(GW_NS, ["sh", "-c", "command -v iptables"], check=False).returncode == 0:
        return "iptables"
    return "none"


def nft_rules() -> list[dict[str, Any]]:
    """Правила таблицы inet awgtest в JSON-виде (nft -j list table)."""
    cp = ns_sh(GW_NS, ["nft", "-j", "list", "table", "inet", NFT_TABLE], check=False)
    if cp.returncode != 0:
        raise StandError(f"nft list table inet {NFT_TABLE} вернул {cp.returncode}: {cp.stderr.strip()}")
    payload = json.loads(cp.stdout)
    return [obj["rule"] for obj in payload.get("nftables", []) if "rule" in obj]


def _nft_rule_counter(rule: dict[str, Any]) -> int | None:
    for expr in rule.get("expr", []):
        if isinstance(expr, dict) and "counter" in expr:
            return int(expr["counter"].get("packets", 0))
    return None


def fw_counters() -> dict[str, Any]:
    """Счётчики правил: {fw, mark_out, mark_fwd, killswitch_drop} (None — недоступно/нет правила)."""
    kind = fw_kind()
    out: dict[str, Any] = {"fw": kind, "mark_out": None, "mark_fwd": None, "killswitch_drop": None}
    if kind == "nft":
        for rule in nft_rules():
            comment = str(rule.get("comment", ""))
            counter = _nft_rule_counter(rule)
            if comment == C_MARK_OUT:
                out["mark_out"] = counter
            elif comment == C_MARK_FWD:
                out["mark_fwd"] = counter
            elif comment == C_DROP:
                out["killswitch_drop"] = counter
    elif kind == "iptables":
        out["mark_out"] = _iptables_counter("-t", "mangle", "-L", IPT_CHAIN_MARK)
        out["mark_fwd"] = _iptables_counter("-t", "mangle", "-L", IPT_CHAIN_FWD)
        out["killswitch_drop"] = _iptables_counter("-t", "filter", "-L", IPT_CHAIN_KS)
    return out


def _iptables_counter(*args: str) -> int | None:
    """Пакеты первого правила в цепочке iptables (в alpine iptables — nft-backend)."""
    cp = ns_sh(GW_NS, ["iptables", *args, "-v", "-n", "-x"], check=False)
    if cp.returncode != 0:
        return None
    for line in cp.stdout.splitlines():
        fields = line.split()
        if len(fields) >= 2 and fields[0].isdigit() and fields[1].isdigit():
            return int(fields[0])
    return None


def iptables_rules() -> str:
    parts = []
    for args in (("-t", "mangle", "-S", IPT_CHAIN_MARK), ("-t", "filter", "-S", IPT_CHAIN_KS)):
        cp = ns_sh(GW_NS, ["iptables", *args], check=False)
        parts.append(cp.stdout.strip() or cp.stderr.strip())
    return "\n".join(parts)


# --- артефакты процесса (логи мока и детектора утечки) -----------------------
def read_jsonl(path: pathlib.Path) -> list[dict[str, Any]]:
    if not path.exists():
        return []
    events: list[dict[str, Any]] = []
    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            events.append(json.loads(line))
        except json.JSONDecodeError:
            continue
    return events


def wait_ready(path: pathlib.Path, timeout: float = 10.0, what: str = "процесс") -> None:
    deadline = time.time() + timeout
    while time.time() < deadline:
        if path.exists():
            return
        time.sleep(0.1)
    raise StandError(f"{what} не поднялся за {timeout}с (нет файла {path})")


def wait_jsonl(path: pathlib.Path, predicate, timeout: float = 5.0) -> dict[str, Any] | None:
    deadline = time.time() + timeout
    while time.time() < deadline:
        for event in read_jsonl(path):
            if predicate(event):
                return event
        time.sleep(0.1)
    return None


# --- управление моком --------------------------------------------------------
def set_mock_mode(mode: str) -> None:
    if mode not in ("respond", "drop"):
        raise ValueError(f"неизвестный режим мока: {mode}")
    MOCK_MODE_FILE.parent.mkdir(parents=True, exist_ok=True)
    MOCK_MODE_FILE.write_text(mode + "\n", encoding="utf-8")
    time.sleep(0.3)  # мок читает режим на каждом пакете


# --- трафик ------------------------------------------------------------------
def probe(
    src_ns: str,
    dst: str,
    port: int,
    *,
    expect_timeout: bool = False,
    raw: str | None = None,
    timeout: float = 2.0,
) -> dict[str, Any]:
    """Отправить тестовый трафик из netns и вернуть разобранный JSON-ответ probe.py."""
    args = ["python3", str(PROBE), "--dst", dst, "--port", str(port), "--timeout", str(timeout)]
    if raw is not None:
        args += ["--raw", raw]
    if expect_timeout:
        args.append("--expect-timeout")
    cp = ns_sh(src_ns, args, check=False, timeout=timeout + 10)
    line = cp.stdout.strip().splitlines()[-1] if cp.stdout.strip() else "{}"
    try:
        result = json.loads(line)
    except json.JSONDecodeError as exc:  # pragma: no cover - диагностика
        raise StandError(f"probe.py вернул не-JSON ({exc}): stdout={cp.stdout!r} stderr={cp.stderr!r}") from exc
    result["exit_code"] = cp.returncode
    return result


def mock_events() -> list[dict[str, Any]]:
    return read_jsonl(MOCK_LOG)


def leak_events() -> list[dict[str, Any]]:
    return [e for e in read_jsonl(LEAK_LOG) if e.get("event") == "leak"]