lab/labsim/labsim-matrix.py

#!/usr/bin/env python3
"""Full-mesh connectivity matrix for the labsim VLANs.

Probes every VLAN VM from every other VLAN VM (ICMP + TCP/22 + TCP/80) and
prints a grid. Use --watch to keep it live: cells that changed since the last
sweep are highlighted, so adding or removing a VyOS firewall rule shows up
within one refresh.

Deliberately dependency-free on the guests: the probe runs with python3, which
is already installed there (cloud-init needs it), so nothing has to be
installed on VMs that have no internet.

  ./labsim-matrix.py                 # one sweep
  ./labsim-matrix.py --watch         # live, refresh every 5s
  ./labsim-matrix.py --watch 2       # live, every 2s
  ./labsim-matrix.py --proto icmp    # single protocol
  ./labsim-matrix.py --json          # machine-readable
"""
from __future__ import annotations

import argparse
import concurrent.futures
import json
import os
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
CONF = os.path.join(HERE, "vlans.conf")

GREEN, RED, GREY, YELLOW, BOLD, RESET = (
    "\033[0;32m", "\033[0;31m", "\033[0;90m", "\033[1;33m", "\033[1m", "\033[0m")

PROTOS = ("icmp", "tcp22", "tcp80")

# Runs ON the guest. Keep it stdlib-only and quick — a hung probe delays the
# whole sweep, so every check is hard-bounded by a timeout.
PROBE = r'''
import json, re, socket, subprocess, sys
targets = json.load(sys.stdin)
out = {}
for name, ip in targets.items():
    res = {}
    try:
        p = subprocess.run(["ping", "-c", "1", "-W", "1", ip],
                           stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, timeout=4)
        res["icmp"] = p.returncode == 0
        # RTT as well as pass/fail: a path that is up but slow is a different
        # problem from one that is down, and the grid alone cannot show it.
        res["rtt_ms"] = None
        if res["icmp"]:
            m = re.search(r"time[=<]\s*([0-9.]+)\s*ms", p.stdout.decode("utf-8", "replace"))
            if m:
                res["rtt_ms"] = float(m.group(1))
    except Exception:
        res["icmp"] = False
        res["rtt_ms"] = None
    for port in (22, 80):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(1.5)
        try:
            s.connect((ip, port)); res["tcp%d" % port] = True
        except Exception:
            res["tcp%d" % port] = False
        finally:
            try: s.close()
            except Exception: pass
    out[name] = res
print(json.dumps(out))
'''


def load_vlans() -> list[dict]:
    vlans = []
    with open(CONF) as fh:
        for line in fh:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            vid, name, prefix, real = line.split(":", 3)
            vlans.append({"vid": vid, "name": name, "ip": f"{prefix}.10",
                          "label": f"{vid}:{name}", "real": real})
    return vlans


def probe_from(src: dict, targets: list[dict], timeout: int) -> tuple[str, dict]:
    """SSH once into src and probe every target from there."""
    payload = json.dumps({t["label"]: t["ip"] for t in targets if t["label"] != src["label"]})
    cmd = [
        "ssh", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null",
        "-o", "BatchMode=yes", "-o", "ConnectTimeout=5", "-o", "LogLevel=ERROR",
        f"alpine@{src['ip']}", "python3", "-",
    ]
    try:
        # The probe script goes on stdin, the target list follows it — the guest
        # reads the script from argv-less stdin, so send both in one stream.
        proc = subprocess.run(
            cmd, input=PROBE.replace("json.load(sys.stdin)", f"json.loads({payload!r})"),
            capture_output=True, text=True, timeout=timeout)
        if proc.returncode != 0:
            return src["label"], {"__error__": (proc.stderr or "ssh failed").strip()[:60]}
        return src["label"], json.loads(proc.stdout)
    except subprocess.TimeoutExpired:
        return src["label"], {"__error__": "probe timed out"}
    except Exception as exc:  # noqa: BLE001 - report, never crash the sweep
        return src["label"], {"__error__": f"{type(exc).__name__}: {exc}"[:60]}


def sweep(vlans: list[dict], timeout: int) -> dict:
    results: dict = {}
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(vlans)) as pool:
        futures = [pool.submit(probe_from, v, vlans, timeout) for v in vlans]
        for fut in concurrent.futures.as_completed(futures):
            label, data = fut.result()
            results[label] = data
    return results


def cell(ok: bool | None, changed: bool) -> str:
    if ok is None:
        return f"{GREY} ·  {RESET}"
    mark = "ok  " if ok else "--  "
    colour = GREEN if ok else RED
    if changed:
        return f"{YELLOW}{BOLD}{'OK*' if ok else 'XX*':<4}{RESET}"
    return f"{colour}{mark}{RESET}"


def render(vlans: list[dict], results: dict, prev: dict | None, protos: tuple[str, ...]) -> None:
    labels = [v["label"] for v in vlans]
    width = max(len(x) for x in labels) + 2

    for proto in protos:
        print(f"\n{BOLD}{proto.upper()}{RESET}   (rows = source, columns = destination)")
        header = " " * width + "".join(f"{lbl:<{width}}" for lbl in labels)
        print(f"{GREY}{header}{RESET}")

        for src in vlans:
            row = f"{src['label']:<{width}}"
            data = results.get(src["label"], {})
            if "__error__" in data:
                print(row + f"{RED}{data['__error__']}{RESET}")
                continue
            for dst in vlans:
                if dst["label"] == src["label"]:
                    row += f"{GREY}{'·':<{width}}{RESET}"
                    continue
                ok = data.get(dst["label"], {}).get(proto)
                was = (prev or {}).get(src["label"], {}).get(dst["label"], {}).get(proto)
                changed = prev is not None and was is not None and was != ok
                txt = cell(ok, changed)
                row += txt + " " * (width - 4)
            print(row)

    reach = sum(1 for s in results.values() if "__error__" not in s
                for d in s.values() for p in protos if d.get(p) is True)
    total = sum(1 for s in results.values() if "__error__" not in s
                for _d in s.values() for _p in protos)
    print(f"\n  reachable: {reach}/{total}   "
          f"{GREEN}ok{RESET}=allowed  {RED}--{RESET}=blocked/no route  "
          f"{YELLOW}*{RESET}=changed since last sweep")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--watch", nargs="?", const=5, type=int, metavar="SECONDS",
                    help="refresh continuously (default every 5s)")
    ap.add_argument("--proto", choices=PROTOS, help="only this protocol")
    ap.add_argument("--json", action="store_true", help="emit raw JSON and exit")
    ap.add_argument("--timeout", type=int, default=30, help="per-host probe timeout")
    args = ap.parse_args()

    vlans = load_vlans()
    protos = (args.proto,) if args.proto else PROTOS

    if args.json:
        print(json.dumps(sweep(vlans, args.timeout), indent=2))
        return 0

    prev = None
    while True:
        started = time.time()
        results = sweep(vlans, args.timeout)
        if args.watch:
            os.system("clear")
            print(f"{BOLD}labsim connectivity matrix{RESET}   "
                  f"{time.strftime('%H:%M:%S')}   (refresh {args.watch}s, Ctrl-C to stop)")
        render(vlans, results, prev, protos)
        if not args.watch:
            return 0
        prev = results
        time.sleep(max(0.0, args.watch - (time.time() - started)))


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print()
        sys.exit(130)