systemdに全部賭けろ（第三部）

#!/usr/bin/env python

from asyncio import Event, create_task, run
from ipaddress import IPv6Address
from json import loads
from time import sleep
from typing import TypedDict

from requests import exceptions, post
from sdbus import (
    DbusInterfaceCommonAsync,
    dbus_method_async,
    dbus_signal_async,
    sd_bus_open_system,
    set_default_bus,
)

SYSTEMD_IF = "org.freedesktop.systemd1"
SYSTEMD_NODE_PATH = "/org/freedesktop/systemd1"
SYSTEMD_MGR_IF = "org.freedesktop.systemd1.Manager"

NETWORKD_IF = "org.freedesktop.network1"
NETWORKD_NODE_PATH = "/org/freedesktop/network1"
NETWORKD_MGR_IF = "org.freedesktop.network1.Manager"
NETWORKD_LINK_IF = "org.freedesktop.network1.Link"
NETWORKD_DHCP6C_IF = "org.freedesktop.network1.DHCPv6Client"

DHCP6C_ETH = "wan0"
HGW_ETH = "hgw"

KEA_API = "http://127.0.0.1:8000"
KEA_SERVICE = "kea-dhcp6.service"
KEA_MAX_RETRIES = 10
KEA_RETRY_DELAY = 5


class Networkd(DbusInterfaceCommonAsync, interface_name=NETWORKD_MGR_IF):
    def __init__(self) -> None:
        super().__init__()
        self._proxify(NETWORKD_IF, NETWORKD_NODE_PATH)

    @dbus_method_async("s", "io")
    async def get_link_by_name(self, name: str) -> tuple[int, str]:
        raise NotImplementedError


class NetworkdLink(DbusInterfaceCommonAsync, interface_name=NETWORKD_LINK_IF):
    def __init__(self, object_path: str) -> None:
        super().__init__()
        self._proxify(NETWORKD_IF, object_path)

    @dbus_method_async(result_signature="s")
    async def describe(self) -> str:
        raise NotImplementedError


class Systemd(DbusInterfaceCommonAsync, interface_name=SYSTEMD_MGR_IF):
    def __init__(self, object_path: str) -> None:
        super().__init__()
        self._proxify(SYSTEMD_IF, SYSTEMD_NODE_PATH)

    @dbus_signal_async(signal_signature="uoss")
    def job_removed(self) -> tuple[int, str, str, str]:
        raise NotImplementedError


class VendorSpecificOption(TypedDict):
    EnterpriseId: int
    SubOptionCode: int
    SubOptionData: str


class Dhcp6(TypedDict):
    prefix: IPv6Address
    dns: str
    ntp: str
    sip: str
    domains: str
    opts: list[VendorSpecificOption]


class Dhcp6State:
    def __init__(self):
        self.state: Dhcp6 | None = None

    def get(self) -> Dhcp6 | None:
        return self.state

    async def set(self, linkpath: str):
        dhcp6 = loads(await NetworkdLink(linkpath).describe())
        dns, ntp, sip = [
            ", ".join([IPv6Address(bytes(e["Address"])).compressed for e in dhcp6[k]])
            for k in ["DNS", "NTP", "SIP"]
        ]
        self.state = {
            "prefix": IPv6Address(bytes(dhcp6["DHCPv6Client"]["Prefixes"][0]["Prefix"])),
            "dns": dns,
            "ntp": ntp,
            "sip": sip,
            "domains": ", ".join([f"{e['Domain']}." for e in dhcp6["SearchDomains"]]),
            "opts": dhcp6["DHCPv6Client"]["VendorSpecificOptions"],
        }

        return self.state


def update_kea_config(dhcp6: Dhcp6):
    def cmd(cmd: str, config: dict[str, str] | None):
        return post(
            KEA_API,
            json={"command": cmd, "service": ["dhcp6"], "arguments": {"Dhcp6": config}},
            timeout=3,
        )

    def get_config():
        for i in range(KEA_MAX_RETRIES):
            try:
                resp = cmd("config-get", None)
                json = resp.json()[0]
                if resp.status_code == 200 and json["result"] == 0:
                    return resp.json()[0]["arguments"]["Dhcp6"]
                else:
                    raise exceptions.ConnectionError
            except exceptions.ConnectionError:
                print(
                    f"Kea API not reachable, retrying in {KEA_RETRY_DELAY} seconds... ({i + 1}/{KEA_MAX_RETRIES})"
                )
                sleep(KEA_RETRY_DELAY)
        print("Kea API is not reachable after maximum retries.")
        return None

    config = get_config()

    if config is None:
        return

    config["subnet6"] = [
        {
            "id": 1,
            "subnet": f"{dhcp6['prefix'] + (0xEF << 64)}/64",
            "interface": HGW_ETH,
            "pools": [{"pool": f"{dhcp6['prefix'] + (0xEF << 64)}/64"}],
            "pd-pools": [
                {
                    "prefix": f"{dhcp6['prefix'] + (0xF0 << 64)}",
                    "prefix-len": 60,
                    "delegated-len": 60,
                }
            ],
        }
    ]

    config["option-def"] = [
        {
            "name": f"ntt-{opt['SubOptionCode']}",
            "code": opt["SubOptionCode"],
            "type": "binary",
            "space": f"vendor-{opt['EnterpriseId']}",
        }
        for opt in dhcp6["opts"]
    ]

    config["option-data"] = [
        {
            "code": opt["SubOptionCode"],
            "always-send": True,
            "space": f"vendor-{opt['EnterpriseId']}",
            "csv-format": False,
            "data": opt["SubOptionData"],
        }
        for opt in dhcp6["opts"]
    ] + [
        {
            "name": "vendor-opts",
            "data": f"{dhcp6['opts'][0]['EnterpriseId']}",
        },
        {"name": "sip-server-addr", "data": dhcp6["sip"]},
        {"name": "dns-servers", "data": dhcp6["dns"]},
        {"name": "domain-search", "data": dhcp6["domains"]},
        {"name": "sntp-servers", "data": dhcp6["ntp"]},
    ]

    test_config = cmd("config-test", config).json()[0]

    print(f"Kea: {test_config['text']}")

    if test_config["result"] != 0:
        return

    set_config = cmd("config-set", config).json()[0]

    if set_config["result"] != 0:
        print(f"Error: {set_config['text']}")
        return

    print(f"Kea: {set_config['text']} Hash: {set_config['arguments']['hash']}")


async def observe_dhcp6_lease(dhcp6_state: Dhcp6State):
    async for linkpath, (
        iface,
        props,
        _,
    ) in DbusInterfaceCommonAsync.properties_changed.catch_anywhere(NETWORKD_IF):
        if iface != NETWORKD_DHCP6C_IF or props["State"][1] != "bound":
            continue

        print(f"DHCPv6 lease acquired on {linkpath}, updating Kea configuration...")

        dhcp6 = await dhcp6_state.set(linkpath)
        update_kea_config(dhcp6)


async def observe_kea_restart(dhcp6_state: Dhcp6State):
    async for _, (
        _,
        _,
        service,
        signal,
    ) in Systemd.job_removed.catch_anywhere(SYSTEMD_IF):
        if service != KEA_SERVICE or signal != "done":
            continue

        print(f"{KEA_SERVICE} restarted, updating Kea configuration...")

        dhcp6 = dhcp6_state.get()

        if dhcp6 is None:
            print("No DHCPv6 state available, fetching from Networkd...")
            _, linkpath = await Networkd().get_link_by_name(DHCP6C_ETH)
            print(f"Link path: {linkpath}")
            dhcp6 = await dhcp6_state.set(linkpath)

        update_kea_config(dhcp6)


async def main():
    set_default_bus(sd_bus_open_system())

    print("Start DHCPv6 proxying...")

    dhcp6_state_manager = Dhcp6State()

    task1 = create_task(observe_dhcp6_lease(dhcp6_state_manager))
    task2 = create_task(observe_kea_restart(dhcp6_state_manager))

    await task1
    await task2

    _ = await Event().wait()


if __name__ == "__main__":
    try:
        run(main())
    except KeyboardInterrupt:
        print("\nExiting...")
    except Exception as e:
        print(f"Error: {e}")