#!/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")
    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:
        print(f"\ndry run -- nothing written. Re-run with --apply to create {len(plan)}.")
        return 0

    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'))}")
    return 0 if fail == 0 and verified == len(plan) else 1


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