utun2/tests/stream_netns/run.py

#!/usr/bin/env python3
"""Запуск стенда: главный namespace используется только для чтения снимка."""

import argparse
import json
import os
from pathlib import Path
import shutil
import signal
import subprocess
import sys
import time

from network import Network
from probe import Probes
from topology import manifest


def host_snapshot():
    result = {}
    for item, arguments in (("addresses", ["ip", "-j", "addr"]), ("routes", ["ip", "-j", "route", "show", "table", "all"]),
                             ("rules", ["ip", "-j", "rule"]), ("qdisc", ["tc", "-j", "qdisc"])):
        result[item] = json.loads(subprocess.check_output(arguments, text=True))
    return result


def inside(args):
    def interrupted(signum, _frame):
        raise KeyboardInterrupt(f"signal {signum}")

    signal.signal(signal.SIGTERM, interrupted)
    signal.signal(signal.SIGINT, interrupted)
    directory = Path(args.output)
    network = Network(directory, args.seed, args.host_netns)
    result = {"pass": False}
    probes = Probes(network)
    try:
        print("Creating 25 isolated network namespaces (16 nodes, 3 routers, 6 NAT gateways)", flush=True)
        network.build()
        network.snapshot("initial")
        print("Calibrating WAN RTT and all 10 NAT clients", flush=True)
        probes.start()
        calibration = probes.calibrate()
        (directory / "calibration.json").write_text(json.dumps(calibration, indent=2) + "\n")
        probes.close()
        print("Network calibration passed", flush=True)
        result["pass"] = True
    except BaseException as error:
        result["error"] = str(error)
        print(f"FAIL: {type(error).__name__}: {error}", file=sys.stderr, flush=True)
    finally:
        probes.close()
        try:
            network.snapshot("final")
        except Exception as error:
            result["snapshot_error"] = str(error)
        network.close()
        (directory / "result.json").write_text(json.dumps(result, indent=2) + "\n")
    return 0 if result["pass"] else 1


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", help="New directory for logs and results")
    parser.add_argument("--seed", type=int, default=20261011)
    parser.add_argument("--network-only", action="store_true", help="Only topology, WAN and NAT calibration")
    parser.add_argument("--inside", action="store_true", help=argparse.SUPPRESS)
    parser.add_argument("--host-netns", help=argparse.SUPPRESS)
    args = parser.parse_args()
    if args.inside:
        return inside(args)
    for program in ("ip", "tc", "nft", "ethtool", "sysctl", "unshare", "nsenter"):
        if not shutil.which(program):
            parser.error(f"Required tool is missing: {program}")
    directory = Path(args.output or Path(__file__).parent / "runs" / time.strftime("%Y%m%d-%H%M%S")).resolve()
    directory.mkdir(parents=True, exist_ok=False)
    (directory / "topology.json").write_text(json.dumps(manifest(args.seed), indent=2) + "\n")
    original_namespace = os.readlink("/proc/self/ns/net")
    before = host_snapshot()
    (directory / "host-before.json").write_text(json.dumps(before, indent=2) + "\n")
    command = ["unshare", "--user", "--map-root-user", "--keep-caps", "--mount", "--net", "--pid", "--fork",
               "--mount-proc", "--propagation", "private", "--kill-child", "--", sys.executable,
               str(Path(__file__).resolve()), "--inside", "--host-netns", original_namespace,
               "--output", str(directory), "--seed", str(args.seed)]
    if args.network_only:
        command.append("--network-only")
    print(f"Artifacts: {directory}", flush=True)
    child = subprocess.Popen(command)
    try:
        status = child.wait()
    except KeyboardInterrupt:
        child.terminate()
        try:
            status = child.wait(timeout=10)
        except subprocess.TimeoutExpired:
            child.kill()
            status = child.wait()
    after = host_snapshot()
    (directory / "host-after.json").write_text(json.dumps(after, indent=2) + "\n")
    unchanged = before == after and os.readlink("/proc/self/ns/net") == original_namespace
    (directory / "host-check.json").write_text(json.dumps({"unchanged": unchanged}, indent=2) + "\n")
    if not unchanged:
        print("FAIL: main network namespace snapshot changed during the run", file=sys.stderr)
        return 1
    print(f"Main network namespace unchanged; exit={status}", flush=True)
    return status


if __name__ == "__main__":
    sys.exit(main())