lab/migration/unifi-reserve-all.py

#!/usr/bin/env python3
"""Reserve every active client at the address it already has.

Why this exists: kea does not inherit UniFi's lease database. At cutover it
starts with an empty view of who holds what, so it can hand an address that is
currently in use to a different device. Reservations are what carry "this
device has this address" across the switch, because they live in config rather
than in lease state.

Dry run by default -- this writes to the live controller, and 40-odd writes is
not something to trigger by accident.

  ./unifi-reserve-all.py                  # show the plan, change nothing
  ./unifi-reserve-all.py --apply          # write them
  ./unifi-reserve-all.py --skip-random    # omit randomised/private MACs

Only clients on networks that actually run DHCP are considered, which
automatically excludes WAN transit VLANs where a reservation is meaningless.
Anything already reserved is left alone, and an address already reserved to a
different MAC is reported and skipped rather than stolen.
"""
from __future__ import annotations

import argparse
import ipaddress
import sys

import _unifi

# The VRRP virtual addresses, read from `show configuration commands` on
# vyos001. UniFi sees these as ordinary client addresses because the firewalls'
# bond MACs answer for them, and their reported IP flips between the real
# interface address and the VIP. Reserving one would put a DHCP reservation on
# the gateway address itself.
VIPS = {
    "192.168.1.254", "192.168.9.254", "192.168.3.254",
    "10.0.9.254", "10.0.1.254", "192.168.2.254",
}

# Every MAC the two firewalls own (bond0/eth0/eth1 share one, eth2 and eth3
# have their own). These interfaces are statically configured routers, not DHCP
# clients -- except eth2, which is deliberately reserved and already handled.
ROUTER_MACS = {
    "64:62:66:25:96:45", "64:62:66:25:96:46", "64:62:66:25:96:48",  # vyos001
    "64:62:66:25:96:51", "64:62:66:25:96:52", "64:62:66:25:96:54",  # vyos002
}


def is_random_mac(mac: str) -> bool:
    """Locally-administered bit set => a privacy/randomised MAC.

    Worth calling out: such a device re-randomises periodically, so the
    reservation stops matching it and becomes dead config. Harmless, but it
    will never do what it looks like it does.
    """
    try:
        return bool(int(mac.split(":")[0], 16) & 0x02)
    except (ValueError, IndexError):
        return False


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--apply", action="store_true", help="actually write (default: dry run)")
    ap.add_argument("--skip-random", action="store_true", help="omit randomised MACs")
    ap.add_argument("--plan", default="reservation-plan.tsv",
                    help="dry run WRITES this file; --apply READS it and applies "
                         "exactly what it contains")
    args = ap.parse_args()

    opener, base, site = _unifi.client()
    users = _unifi.get(opener, base, site, "rest/user")
    nets = _unifi.get(opener, base, site, "rest/networkconf")
    sta = _unifi.get(opener, base, site, "stat/sta")
    for blob in (users, nets, sta):
        if isinstance(blob, dict):
            print(f"error reading controller: {blob['__error__']}", file=sys.stderr)
            return 1

    # Only networks that serve DHCP: a reservation on a WAN transit VLAN means
    # nothing, and those are exactly the ones without dhcpd_enabled.
    serving = []
    for n in nets:
        if not (n.get("dhcpd_enabled") and n.get("ip_subnet")):
            continue
        try:
            serving.append((ipaddress.ip_interface(n["ip_subnet"]).network, n))
        except ValueError:
            continue

    by_mac = {(u.get("mac") or "").lower(): u for u in users}
    taken = {u.get("fixed_ip"): (u.get("mac") or "").lower()
             for u in users if u.get("use_fixedip")}

    plan, skipped = [], []
    for c in sta:
        mac, ip = (c.get("mac") or "").lower(), c.get("ip")
        if not mac or not ip:
            continue
        label = c.get("name") or c.get("hostname") or "?"
        user = by_mac.get(mac)

        if ip in VIPS:
            skipped.append((ip, mac, label, "VRRP virtual address - not a client"))
            continue
        if mac in ROUTER_MACS:
            skipped.append((ip, mac, label, "firewall's own interface - statically configured"))
            continue

        net = next((n for netw, n in serving if ipaddress.ip_address(ip) in netw), None)
        if net is None:
            skipped.append((ip, mac, label, "not on a DHCP-serving network"))
            continue
        # The gateway is the router, not a lease.
        if ip == str(ipaddress.ip_interface(net["ip_subnet"]).ip):
            skipped.append((ip, mac, label, "network gateway address"))
            continue
        if user is None:
            skipped.append((ip, mac, label, "not a known client on the controller"))
            continue
        if user.get("use_fixedip"):
            if user.get("fixed_ip") != ip:
                skipped.append((ip, mac, label,
                                f"already reserved at {user.get('fixed_ip')} - left alone"))
            continue
        if ip in taken and taken[ip] != mac:
            skipped.append((ip, mac, label,
                            f"address already reserved to {taken[ip]}"))
            continue
        if args.skip_random and is_random_mac(mac):
            skipped.append((ip, mac, label, "randomised MAC (--skip-random)"))
            continue
        plan.append((ip, mac, label, user, net))

    # Generic safety net: if two MACs report the same current address, at most
    # one of them can legitimately keep it and we cannot tell which. Drop both
    # and say so -- this is exactly how the VRRP VIPs first showed up.
    counts: dict[str, int] = {}
    for ip, *_ in plan:
        counts[ip] = counts.get(ip, 0) + 1
    contested = {ip for ip, n in counts.items() if n > 1}
    if contested:
        for ip, mac, label, _u, _n in [p for p in plan if p[0] in contested]:
            skipped.append((ip, mac, label, "address claimed by more than one MAC"))
        plan = [p for p in plan if p[0] not in contested]

    plan.sort(key=lambda r: ipaddress.ip_address(r[0]))

    print(f"=== plan: {len(plan)} new reservation(s) ===")
    for ip, mac, label, _u, net in plan:
        flag = "  [randomised MAC]" if is_random_mac(mac) else ""
        print(f"  {ip:16} {mac:18} {label[:28]:28} {net.get('name')}{flag}")
    if skipped:
        print(f"\n=== skipped ({len(skipped)}) ===")
        for ip, mac, label, why in sorted(skipped):
            print(f"  {ip:16} {mac:18} {label[:24]:24} {why}")

    n_rand = sum(1 for p in plan if is_random_mac(p[1]))
    if n_rand:
        print(f"\nnote: {n_rand} of these use randomised MACs. The reservation "
              f"stops matching once the device re-randomises.")

    if not args.apply:
        with open(args.plan, "w") as fh:
            for ip, mac, label, _u, _n in plan:
                fh.write(f"{mac}\t{ip}\t{label}\n")
        print(f"\ndry run -- nothing written to the controller.")
        print(f"plan saved to {args.plan}; re-run with --apply to apply exactly that.")
        return 0

    # Apply the plan that was REVIEWED, not one recomputed now.
    #
    # This cost a k8s node an outage. The apply used to re-read stat/sta, and a
    # client that renewed between the dry run and the apply got pinned to
    # whatever transient address it happened to hold at that instant -- worker1
    # was reviewed at .13 and written as .242. A plan you looked at and a plan
    # that gets applied must be the same object.
    try:
        with open(args.plan) as fh:
            reviewed = {}
            for line in fh:
                parts = line.rstrip("\n").split("\t")
                if len(parts) >= 2:
                    reviewed[parts[0].lower()] = parts[1]
    except OSError:
        print(f"no plan at {args.plan}. Run without --apply first and review it.",
              file=sys.stderr)
        return 1

    drifted = [(ip, mac) for ip, mac, _l, _u, _n in plan
               if mac in reviewed and reviewed[mac] != ip]
    for ip, mac in drifted:
        print(f"  note: {mac} now reports {ip}, plan says {reviewed[mac]} -- "
              f"applying the plan", file=sys.stderr)

    plan = [(reviewed[mac], mac, label, user, net)
            for ip, mac, label, user, net in plan if mac in reviewed]
    print(f"applying {len(plan)} reservation(s) from {args.plan}")

    print()
    ok = fail = 0
    for ip, mac, label, user, net in plan:
        res = _unifi.put(opener, base, site, f"rest/user/{user['_id']}",
                         {"use_fixedip": True, "fixed_ip": ip, "network_id": net["_id"]})
        if isinstance(res, dict):
            print(f"  FAILED {ip:16} {mac}  {res['__error__'][:70]}")
            fail += 1
        else:
            ok += 1

    # Read back rather than trusting the write responses.
    after = _unifi.get(opener, base, site, "rest/user")
    live = {(u.get("mac") or "").lower() for u in after if u.get("use_fixedip")}
    verified = sum(1 for _ip, mac, _l, _u, _n in plan if mac in live)
    print(f"\nwrote {ok}, failed {fail}, verified live {verified}/{len(plan)}")
    print(f"total reservations on the controller now: "
          f"{sum(1 for u in after if u.get('use_fixedip'))}")

    # Writing the controller is only half the job. The gateway applies config
    # on its own schedule, and a device running config from before these
    # changes will hand out addresses that disagree with what the controller
    # shows -- which is how a k8s node ended up unable to get any lease at all
    # while the controller looked perfectly correct.
    print("\nThe controller now disagrees with what the gateway is running.")
    print("Push it to the device and wait for state to return to 'connected':")
    print("  python3 -c \"import _unifi; o,b,s=_unifi.client(); "
          "print(_unifi.post(o,b,s,'cmd/devmgr',"
          "{'cmd':'force-provision','mac':'<gateway-mac>'}))\"")
    print("Then verify a real DISCOVER is answered before trusting it:")
    print("  ssh vyos@<fw> 'sudo nmap --script broadcast-dhcp-discover -e eth2 "
          "--script-args broadcast-dhcp-discover.mac=<client-mac>'")
    return 0 if fail == 0 and verified == len(plan) else 1


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