dotfiles/nixos/tests/router.py

from test_driver.machine import Machine, StartCommand, NixStartScript
from test_driver.logger import TerminalLogger
from pathlib import Path
from typing import Any
import os
import sys
import socket
import subprocess
import pty
import time

hela_system = os.getenv("HELA_SYSTEM")
client_system = os.getenv("CLIENT_SYSTEM")
jormungandr_system = os.getenv("JORMUNGANDR_SYSTEM")
interactive = os.getenv("ROUTER_INTERACTIVE")

class FileLogger(TerminalLogger):
    @staticmethod
    def _eprint(*args: object, **kwargs: Any) -> None:
        return

class NixStartScriptSetsid(NixStartScript):
    def run(
        self,
        state_dir: Path,
        shared_dir: Path,
        monitor_socket_path: Path,
        qmp_socket_path: Path,
        shell_socket_path: Path,
        allow_reboot: bool,
    ) -> subprocess.Popen:
        return subprocess.Popen(
            self.cmd(
                monitor_socket_path, qmp_socket_path, shell_socket_path, allow_reboot
            ),
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            shell=True,
            cwd=state_dir,
            env=self.build_environment(state_dir, shared_dir),
        )

def make_machine(*, system: Path, name: str, out_dir: Path | None = None, tmp_dir: Path | None = None) -> Machine:
    if out_dir is None:
        out_dir = Path(f"{name}/out")
    if tmp_dir is None:
        tmp_dir = Path(f"{name}/tmp")
    out_dir.mkdir(parents=True, exist_ok=True)
    tmp_dir.mkdir(parents=True, exist_ok=True)
    return Machine(
        out_dir = out_dir.absolute(),
        tmp_dir = tmp_dir.absolute(),
        start_command = NixStartScript(system + f"/bin/run-{name}-vm") if not interactive else NixStartScriptSetsid(system + f"/bin/run-{name}-vm"),
        logger = TerminalLogger() if not interactive else FileLogger(),
        name = name,
        keep_vm_state = True,
    )

def vde_switch(control_socket: Path) -> subprocess.Popen:
    control_socket.mkdir(parents = True, exist_ok = True)
    pty_master, pty_slave = pty.openpty()
    return subprocess.Popen(
        ["vde_switch", "-unix", str(control_socket.absolute()), "--dirmode", "0700"],
        stdin=pty_slave,
    )

def vde_plug_unix(path: Path, control_socket: Path, port: int | None = None) -> subprocess.Popen:
    port = f"-p {port}" if port is not None else ""
    return subprocess.Popen(
        ["socat", f"UNIX-LISTEN:{path.absolute()}", f"EXEC:vde_plug {port} {control_socket.absolute()}"],
        stdin = subprocess.PIPE,
        stdout = subprocess.PIPE,
        stderr = subprocess.STDOUT,
    )

out_dir = Path(os.getenv("out"))
def dump_to_out(path: str | Path, content: bytes):
    global out_dir

    if not isinstance(path, Path):
        path = Path(path)

    final_path = out_dir / path
    final_path.parent.mkdir(parents = True, exist_ok = True)

    fh = open(final_path, "w")
    fh.write(content)
    fh.close()

vde_hel_control_path = Path("hel/switch")
vde_hel = vde_switch(vde_hel_control_path)

vde_midgard_control_path = Path("midgard/switch")
vde_midgard = vde_switch(vde_midgard_control_path)

# kpn_link = subprocess.Popen(
#     [
#         "socat",
#         f"UNIX-LISTEN:{Path('midgard/link.jormungandr').absolute()}",
#         f"UNIX-LISTEN:{Path('midgard/link.hela').absolute()}"
#     ]
# )

time.sleep(0.5)

hela = make_machine(system = hela_system, name = "hela")
hela.start()

time.sleep(0.5)

client = make_machine(system = client_system, name = "client")
client.start()

jormungandr = make_machine(system = jormungandr_system, name = "jormungandr")
jormungandr.start()

def finalize(exit_code: int) -> None:
    global hela, client, jormungandr, vde_hel, vde_midgard

    # subprocess.run(["cp", "hela/tmp/vm-state-hela/xchg/tcpdump-dmz-log:0.log", out_dir / "hela-tcpdump-dmz-log:0.log"], check=True)
    # subprocess.run(["cp", "hela/tmp/vm-state-hela/xchg/tcpdump-dmz-log:1.log", out_dir / "hela-tcpdump-dmz-log:1.log"], check=True)

    # subprocess.run(["cp", "jormungandr/tmp/vm-state-jormungandr/xchg/tcpdump-dmz-log:0.log", out_dir / "jormungandr-tcpdump-dmz-log:0.log"], check=True)
    # subprocess.run(["cp", "jormungandr/tmp/vm-state-jormungandr/xchg/tcpdump-dmz-log:1.log", out_dir / "jormungandr-tcpdump-dmz-log:1.log"], check=True)

    hela.shutdown()
    client.shutdown()
    jormungandr.shutdown()

    vde_hel.terminate()
    vde_midgard.terminate()

    vde_hel.wait()
    vde_midgard.wait()

    sys.exit(exit_code)

try:
    if interactive:
        import ptpython

        ptpython.repl.embed(globals(), locals())
    else:
        hela.wait_for_unit("multi-user.target")
        client.wait_for_unit("multi-user.target")
        jormungandr.wait_for_unit("multi-user.target")

        # hela.wait_until_succeeds("ip netns exec dmz bash -c $'[ $(ip -f inet addr show ppp0 | sed -En -e \\'s/.*inet ([0-9.]+).*/\\1/p\\' | wc -l) -gt 0 ]'")

        dump_to_out("client-default-route", client.succeed("ip route"))
        dump_to_out("hela-default-route", hela.succeed("ip route"))
        dump_to_out("hela-dmz-route", hela.succeed("ip netns exec dmz ip route"))

        dump_to_out("client-default-if", client.succeed("ip addr"))
        dump_to_out("hela-default-if", hela.succeed("ip addr"))
        dump_to_out("hela-dmz-if", hela.succeed("ip netns exec dmz ip addr"))

        client.succeed("ping -c 1 10.10.0.1")
        hela.succeed("ping -c 1 10.10.0.2")

        hela.succeed("ip netns exec dmz ping -c 1 10.0.0.1")
        jormungandr.fail("ping -W 1 -c 1 10.67.15.1")

        # client.succeed("ping -c 1 10.0.0.1")

# hela.succeed("systemd-run -u tcpdump-dmz-wan -p StandardOutput=append:/tmp/xchg/tcpdump-dmz-log:0.log -p StandardError=append:/tmp/xchg/tcpdump-dmz-log:1.log ip netns exec dmz tcpdump -i internet")
# jormungandr.succeed("systemd-run -u tcpdump-dmz-wan -p StandardOutput=append:/tmp/xchg/tcpdump-dmz-log:0.log -p StandardError=append:/tmp/xchg/tcpdump-dmz-log:1.log tcpdump -i internet")

except Exception as err:
    print(err)

    finalize(1)

finalize(0)