bird/python/BIRD/Test.py

import asyncio
import ipaddress
import jinja2
import os
import pathlib
import sys
import yaml

sys.path.insert(0, "/home/maria/flock")

from flock.Hypervisor import Hypervisor
from flock.Machine import Machine
from .CLI import CLI, Transport

# TODO: move this to some aux file
class Differs(Exception):
    def __init__(self, a, b, tree):
        self.a = a
        self.b = b
        self.tree = tree

    @classmethod
    def false(cls, a, b, deep, tree):
        if deep:
            raise cls(a, b, tree)
        else:
            return False

def deep_eq(a, b, deep=False):
    if a == b:
        return True

    # Do not iterate over strings
    if type(a) is str and type(b) is str:
        return Differs.false(a, b, deep, tree=[])

    try:
        for k,v in a.items():
            try:
                deep_eq(v, b[k], True)
            except Differs as d:
                d.tree.append(k)
                raise d
            except KeyError:
                return Differs.false(v, None, deep, tree=[k])

        for k in b:
            if not k in a:
                return Differs.false(None, b[k], deep, tree=[k])

    except AttributeError:
        try:
            if len(a) != len(b):
                return Differs.false(len(a), len(b), deep, tree=[len])

            for i in range(len(a)):
                try:
                    deep_eq(a[i], b[i])
                except Differs as d:
                    d.tree.append(i)
                    raise d
        except TypeError:
            return Differs.false(a, b, deep, [])

    return True

class MinimalistTransport(Transport):
    def __init__(self, socket, machine):
        self.sock = socket
        self.machine = machine

    async def send_cmd(self, *args):
        return await self.sock.send_cmd("run_in", self.machine, "./birdc", "-l", *args)

class BIRDBinDir:
    index = {}

    def __init__(self, path):
        self.path = path

        f = [ "bird", "birdc", "birdcl", ]

        self.files = { k: None for k in f }
        self.mod = { k: None for k in f }
        self.loaded = False

    @classmethod
    def get(cls, where):
        w = pathlib.Path(where).absolute()
        try:
            return cls.index[s := str(w)]
        except KeyError:
            cls.index[s] = (b := cls(w))
            return b

    def load(self):
        for bn,v in self.files.items():
            if v is None:
                with open(self.path / bn, "rb") as b:
                    self.files[bn] = b.read()
                self.mod[bn] = (self.path / bn).stat().st_mode

        self.loaded = True

    def copy(self, target):
        if not self.loaded:
            self.load()

        for bn,v in self.files.items():
            if v is not None:
                with open(target / bn, "wb") as b:
                    b.write(v)
                (target / bn).chmod(self.mod[bn])

    def cleanup(self, target):
        for bn in self.files:
            (target / bn).unlink()

default_bindir = BIRDBinDir.get(".")

class BIRDInstance(CLI):
    def __init__(self, mach: Machine, bindir=None, conf=None):
        self.mach = mach
        self.workdir = self.mach.workdir
        self.bindir = BIRDBinDir.get(bindir) if bindir is not None else default_bindir
        self.conf = conf if conf is not None else f"bird_{mach.name}.conf"

        super().__init__(
                transport=MinimalistTransport(
                    socket=mach.hypervisor.control_socket,
                    machine=self.mach.name
                    )
                )

    async def start(self, test):
        self.bindir.copy(self.workdir)

        with (open(self.conf, "r") as s, open(self.workdir / "bird.conf", "w") as f):
            f.write(jinja2.Environment().from_string(s.read()).render(t=test))

        await test.hcom("run_in", self.mach.name, "./bird", "-l")

    async def cleanup(self):
        # Send down command and wait for BIRD to actually finish
        await self.down()
        while (self.workdir / "bird.ctl").exists():
            await asyncio.sleep(0.1)

        # Remove known files
        for f in ("bird.conf", "bird.log"):
            (self.workdir / f).unlink()

        self.bindir.cleanup(self.workdir)

class Test:
    ipv6_prefix = ipaddress.ip_network("2001:db8::/32")
    ipv4_prefix = ipaddress.ip_network("192.0.2.0/24")

    ipv6_link_pxlen = 64
    ipv4_link_pxlen = 28

    SAVE = 1
    CHECK = 2

    show_difs = False

    # 198.51.100.0/24, 203.0.113.0/24

    def __init__(self, name, mode: int):
        self.name = name
        self.hypervisor = Hypervisor(name)
        self.machine_index = {}
        self.mode = mode
        self._started = None
        self._starting = False

        self.ipv6_pxgen = self.ipv6_prefix.subnets(new_prefix=self.ipv6_link_pxlen)
        self.ipv4_pxgen = self.ipv4_prefix.subnets(new_prefix=self.ipv4_link_pxlen)

        self.route_dump_id = 0

    async def hcom(self, *args):
        if self._started is None:
            self._started = asyncio.Future()

        if self._started.done():
            return await self.hypervisor.control_socket.send_cmd(*args)

        if self._starting:
            await self._started
        else:
            self._starting = True
            await self.hypervisor.prepare()
            os.symlink(pathlib.Path("bgp-secondary.log").absolute(), self.hypervisor.basedir / "flock.log")
            await self.hypervisor.start()

            self._started.set_result(True)
            self._starting = False

        return await self.hypervisor.control_socket.send_cmd_early(*args)

    async def machines(self, *names, t: type):
        for n in names:
            if n in self.machine_index:
                raise Exception(f"Machine {n} duplicate")

        info = await asyncio.gather(*[
            self.hcom("machine", name, { "type": "minimalist" })
            for name in names
            ])

        inst = [
                t(mach=Machine.new(
                    name=n,
                    hypervisor=self.hypervisor,
                    **i
                    ),
                  ) for n,i in zip(names, info)
                ]

        for n,i in zip(names, inst):
            self.machine_index[n] = i

        return inst

    async def link(self, name, *machines):
        match len(machines):
            case 0:
                raise Exception("Link with no machines? HOW?!")
            case 1:
                raise NotImplementedError("dummy link")
            case 2:
                linfo = await self.hcom("link", name, {
                    "machines": { m: { "name": name } for m in machines },
                    "ipv6": str(next(self.ipv6_pxgen)),
                    "ipv4": str(next(self.ipv4_pxgen)),
                    })
                for m in machines:
                    for i in ("ipv4", "ipv6"):
                        linfo[m][i] = ipaddress.ip_interface(linfo[m][i])

                return linfo

            case _:
                raise NotImplementedError("virtual bridge")

    async def start(self):
        await asyncio.gather(*[ v.start(test=self) for v in self.machine_index.values() ])

    async def cleanup(self):
        await asyncio.gather(*[ v.cleanup() for v in self.machine_index.values() ])
        await self.hcom("stop", True)

    async def route_dump(self, timeout, name, machines=None, check_timeout=10, check_retry_timeout=0.5):
        # Compile dump ID
        self.route_dump_id += 1
        if name is None:
            name = f"dump-{self.route_dump_id:04d}.yaml"
        else:
            name = f"dump-{self.route_dump_id:04d}-{name}.yaml"

        # Collect machines to dump
        if machines is None:
            machines = self.machine_index.values()
        else:
            machines = [
                    m if isinstance(m, CLI) else self.machine_index[m]
                    for m in machines
                    ]

        # Define the obtainer function
        async def obtain():
            dump = await asyncio.gather(*[
                where.show_route()
                for where in machines
                ])
            for d in dump:
                for t in d["tables"].values():
                    for n in t["networks"].values():
                        for r in n["routes"]:
                            assert("when" in r)
                            r["when"] = True
            return dump

        match self.mode:
            case Test.SAVE:
                await asyncio.sleep(timeout)
                dump = await obtain()
                with open(name, "w") as y:
                    yaml.dump_all(dump, y)
                print(f"{name}\t{self.route_dump_id}\t[ SAVED ]")

            case Test.CHECK:
                with open(name, "r") as y:
                    c = [*yaml.safe_load_all(y)]

                seen = []
                try:
                    async with asyncio.timeout(check_timeout) as to:
                        while True:
                            dump = await obtain()
                            try:
                                deep_eq(c, dump, True)
#                            if deep_eq(c, dump):
                                spent = asyncio.get_running_loop().time() - to.when() + check_timeout
                                print(f"{name}\t{self.route_dump_id}\t[  OOK  ]\t{spent:.6f}s")
                                return True
                            except Differs as d:
                                if self.show_difs:
                                    print(f"Differs at {' -> '.join([str(s) for s in reversed(d.tree)])}: {d.a} != {d.b}")

                            seen.append(dump)
                            await asyncio.sleep(check_retry_timeout)
                except TimeoutError as e:
                    print(f"{name}\t{self.route_dump_id}\t[ BAD  ]")
                    for q in range(len(seen)):
                        with open(f"__result_bad_{q}__{name}", "w") as y:
                            yaml.dump_all(seen[q], y)

                    return False

            case _:
                raise Exception("Invalid test mode")

if __name__ == "__main__":
    name = sys.argv[1]
    mode = Test.CHECK

    if name == "-s":
        name = sys.argv[2]
        mode = Test.SAVE

    p = (pathlib.Path(__file__).parent.parent.parent / "flock" / name).absolute()
    sys.path.insert(0, str(p))

#    if "MAKEFLAGS" in os.environ:
#        print(os.environ["MAKEFLAGS"])

    import test

    os.chdir(p)
    asyncio.run(test.ThisTest(name, mode).run())