#!/usr/bin/env python3
"""Export everything from the UniFi controller that the VyOS cutover must preserve.

The point is that nothing quietly stops working after the switch. That means
capturing not just the networks but every DHCP reservation, every port forward
and every firewall rule — the things nobody remembers configuring until they
break.

Writes one JSON file per endpoint into ./export/ (raw, unmodified — the source
of truth) plus inventory.json, a normalised view used by the VyOS generator.

  ./unifi-export.py                 # export to ./export/
  ./unifi-export.py --out /tmp/x    # elsewhere
  ./unifi-export.py --summary       # print a human summary of what was found

WARNING: the raw export contains secrets (wlanconf holds WiFi passphrases,
setting holds RADIUS/auth material). ./export/ is gitignored — keep it that way.
"""
from __future__ import annotations

import argparse
import ipaddress
import json
import os
import sys

import _unifi

# endpoint -> why it matters for the cutover
ENDPOINTS = {
    "rest/networkconf": "networks: VLANs, subnets, DHCP ranges, DNS, lease time",
    "rest/user":        "known clients — this is where fixed DHCP reservations live",
    "stat/sta":         "currently active clients and their live IPs",
    "rest/firewallrule": "firewall rules",
    "rest/firewallgroup": "address/port groups referenced by rules",
    "rest/portforward": "port forwards (inbound NAT)",
    "rest/routing":     "static routes",
    "rest/dhcpoption":  "custom DHCP options",
    "rest/wlanconf":    "wireless networks (VLAN bindings)",
    "stat/device":      "switches/APs incl. per-port VLAN config",
    "rest/setting":     "controller settings (incl. USG/gateway config)",
    "rest/usergroup":   "bandwidth groups referenced by clients",
    "rest/dynamicdns":  "dynamic DNS",
}


def build_inventory(raw: dict) -> dict:
    """Normalise the parts a migration actually has to reproduce."""
    nets_by_id = {n["_id"]: n for n in raw.get("rest/networkconf", []) if isinstance(n, dict)}

    networks = []
    for n in raw.get("rest/networkconf", []):
        if not isinstance(n, dict):
            continue
        networks.append({
            "id": n.get("_id"),
            "name": n.get("name"),
            "purpose": n.get("purpose"),
            "vlan": n.get("vlan"),
            "vlan_enabled": n.get("vlan_enabled"),
            "subnet": n.get("ip_subnet"),
            "domain_name": n.get("domain_name"),
            "dhcp_enabled": n.get("dhcpd_enabled"),
            "dhcp_start": n.get("dhcpd_start"),
            "dhcp_stop": n.get("dhcpd_stop"),
            "dhcp_lease": n.get("dhcpd_leasetime"),
            "dhcp_dns": [n.get(f"dhcpd_dns_{i}") for i in (1, 2, 3, 4) if n.get(f"dhcpd_dns_{i}")],
            "dhcp_gateway": n.get("dhcpd_gateway") or n.get("dhcpd_gateway_enabled"),
            "dhcp_ntp": [n.get(f"dhcpd_ntp_{i}") for i in (1, 2) if n.get(f"dhcpd_ntp_{i}")],
            "igmp_snooping": n.get("igmp_snooping"),
            "enabled": n.get("enabled", True),
        })

    # ip_subnet is the GATEWAY address with a prefix ("10.0.0.1/23"), not the
    # network address — so derive the real subnet before matching against it.
    subnets = []
    for n in networks:
        if not n["subnet"]:
            continue
        try:
            iface = ipaddress.ip_interface(n["subnet"])
        except ValueError:
            continue
        subnets.append((iface.network, n))

    def resolve_net(ip: str | None, network_id: str | None) -> dict:
        """Which network does this reservation belong to?

        Most reservations here (23 of 31 as of the first export) carry no
        network_id at all — UniFi simply does not bind them. VyOS needs the
        subnet to place a static-mapping, so fall back to containment.
        """
        if network_id and network_id in nets_by_id:
            nid = nets_by_id[network_id]
            return {"name": nid.get("name"), "vlan": nid.get("vlan"), "by": "network_id"}
        if ip:
            try:
                addr = ipaddress.ip_address(ip)
            except ValueError:
                return {"name": None, "vlan": None, "by": "unresolved"}
            for net, meta in subnets:
                if addr in net:
                    return {"name": meta["name"], "vlan": meta["vlan"], "by": "subnet"}
        return {"name": None, "vlan": None, "by": "unresolved"}

    # Fixed reservations: the single most important thing to carry over, and
    # the easiest to lose — nobody has these written down anywhere else.
    reservations = []
    for u in raw.get("rest/user", []):
        if not isinstance(u, dict) or not u.get("use_fixedip"):
            continue
        net = resolve_net(u.get("fixed_ip"), u.get("network_id"))
        reservations.append({
            "mac": (u.get("mac") or "").lower(),
            "ip": u.get("fixed_ip"),
            "name": u.get("name") or u.get("hostname") or "",
            "hostname": u.get("hostname") or "",
            "network_id": u.get("network_id"),
            "network_name": net["name"],
            "network_vlan": net["vlan"],
            "resolved_by": net["by"],
            "note": (u.get("note") or "").strip(),
        })
    reservations.sort(key=lambda r: tuple(int(p) for p in r["ip"].split(".")) if r["ip"] else (0,))

    # Active leases without a reservation: these devices work today by luck of
    # the lease database. After a DHCP server swap they get a NEW address.
    reserved_macs = {r["mac"] for r in reservations}
    dynamic = []
    for c in raw.get("stat/sta", []):
        if not isinstance(c, dict):
            continue
        mac = (c.get("mac") or "").lower()
        if mac in reserved_macs or not c.get("ip"):
            continue
        dynamic.append({
            "mac": mac,
            "ip": c.get("ip"),
            "name": c.get("name") or c.get("hostname") or "",
            "network": c.get("network"),
        })
    dynamic.sort(key=lambda r: tuple(int(p) for p in r["ip"].split(".")) if r["ip"] else (0,))

    port_forwards = [{
        "name": p.get("name"), "enabled": p.get("enabled"),
        "proto": p.get("proto"), "src": p.get("src"),
        "dst_port": p.get("dst_port"), "fwd": p.get("fwd"),
        "fwd_port": p.get("fwd_port"), "log": p.get("log"),
    } for p in raw.get("rest/portforward", []) if isinstance(p, dict)]

    firewall_rules = [{
        "name": r.get("name"), "enabled": r.get("enabled"), "action": r.get("action"),
        "ruleset": r.get("ruleset"), "rule_index": r.get("rule_index"),
        "protocol": r.get("protocol"),
        "src_address": r.get("src_address"), "dst_address": r.get("dst_address"),
        "src_firewallgroup_ids": r.get("src_firewallgroup_ids"),
        "dst_firewallgroup_ids": r.get("dst_firewallgroup_ids"),
        "src_networkconf_id": r.get("src_networkconf_id"),
        "dst_networkconf_id": r.get("dst_networkconf_id"),
    } for r in raw.get("rest/firewallrule", []) if isinstance(r, dict)]

    firewall_groups = [{
        "id": g.get("_id"), "name": g.get("name"),
        "type": g.get("group_type"), "members": g.get("group_members"),
    } for g in raw.get("rest/firewallgroup", []) if isinstance(g, dict)]

    static_routes = [{
        "name": r.get("name"), "enabled": r.get("enabled"),
        "network": r.get("static-route_network"),
        "nexthop": r.get("static-route_nexthop"),
        "distance": r.get("static-route_distance"),
        "type": r.get("static-route_type"),
    } for r in raw.get("rest/routing", []) if isinstance(r, dict)]

    return {
        "networks": networks,
        "reservations": reservations,
        "dynamic_clients": dynamic,
        "port_forwards": port_forwards,
        "firewall_rules": firewall_rules,
        "firewall_groups": firewall_groups,
        "static_routes": static_routes,
        "warnings": find_warnings(networks, reservations),
    }


def find_warnings(networks: list, reservations: list) -> list:
    """Things that are fine under UniFi but bite when rebuilt on VyOS."""
    warns = []

    for n in networks:
        if not n["subnet"]:
            continue
        try:
            iface = ipaddress.ip_interface(n["subnet"])
        except ValueError:
            warns.append({"kind": "bad_subnet", "network": n["name"], "detail": n["subnet"]})
            continue
        # UniFi stores the gateway in ip_subnet. A gateway equal to the network
        # address is legal in a /23 but plenty of tooling rejects it, so it must
        # not be discovered during the cutover window.
        if iface.ip == iface.network.network_address:
            warns.append({
                "kind": "gateway_is_network_address", "network": n["name"],
                "detail": f"gateway {iface.ip} is the network address of {iface.network}",
            })

    # Reservations that sit inside the dynamic pool. UniFi's dhcpd tolerates
    # this; whether VyOS does depends on its DHCP backend, so every one of these
    # is a config that must be proven on the sim before cutover.
    ranges = []
    for n in networks:
        if n["dhcp_enabled"] and n["dhcp_start"] and n["dhcp_stop"]:
            try:
                ranges.append((n["name"], ipaddress.ip_address(n["dhcp_start"]),
                               ipaddress.ip_address(n["dhcp_stop"])))
            except ValueError:
                pass
    inside = []
    for r in reservations:
        if not r["ip"]:
            continue
        try:
            addr = ipaddress.ip_address(r["ip"])
        except ValueError:
            continue
        for name, lo, hi in ranges:
            if lo <= addr <= hi:
                inside.append(f"{r['ip']} ({r['name'] or r['mac']}) in {name} pool")
                break
    if inside:
        warns.append({"kind": "reservation_inside_dhcp_pool",
                      "count": len(inside), "detail": inside})

    unresolved = [f"{r['ip']} {r['mac']} {r['name']}"
                  for r in reservations if r["resolved_by"] == "unresolved"]
    if unresolved:
        warns.append({"kind": "reservation_matches_no_subnet",
                      "count": len(unresolved), "detail": unresolved})
    return warns


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "export"))
    ap.add_argument("--summary", action="store_true")
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)
    opener, base, site = _unifi.client()
    print(f"controller {base}  site {site}")

    raw: dict = {}
    errors = []
    for path, why in ENDPOINTS.items():
        data = _unifi.get(opener, base, site, path)
        if isinstance(data, dict) and "__error__" in data:
            errors.append((path, data["__error__"]))
            print(f"  {path:22} FAILED  {data['__error__'][:50]}")
            continue
        raw[path] = data
        fname = path.replace("/", "_") + ".json"
        with open(os.path.join(args.out, fname), "w") as fh:
            json.dump(data, fh, indent=2, sort_keys=True)
        print(f"  {path:22} {len(data):4d}  -> {fname}")

    inventory = build_inventory(raw)
    with open(os.path.join(args.out, "inventory.json"), "w") as fh:
        json.dump(inventory, fh, indent=2, sort_keys=True)

    print(f"\nwrote {args.out}/inventory.json")
    print(f"  networks           {len(inventory['networks'])}")
    print(f"  DHCP reservations  {len(inventory['reservations'])}")
    print(f"  dynamic clients    {len(inventory['dynamic_clients'])}  (no reservation — see summary)")
    print(f"  port forwards      {len(inventory['port_forwards'])}")
    print(f"  firewall rules     {len(inventory['firewall_rules'])}")
    print(f"  static routes      {len(inventory['static_routes'])}")
    if errors:
        print(f"  endpoints failed   {len(errors)}: {[e[0] for e in errors]}")

    if args.summary:
        print_summary(inventory)
    return 0


def print_summary(inv: dict) -> None:
    print("\n=== networks ===")
    print(f"{'name':22} {'vlan':>5}  {'subnet':20} {'dhcp range':32} lease")
    for n in sorted(inv["networks"], key=lambda x: (x["vlan"] or 0)):
        rng = f"{n['dhcp_start']} - {n['dhcp_stop']}" if n["dhcp_enabled"] else "(dhcp off)"
        print(f"{(n['name'] or '')[:22]:22} {str(n['vlan'] or '-'):>5}  "
              f"{(n['subnet'] or '-'):20} {rng:32} {n['dhcp_lease'] or '-'}")

    print(f"\n=== DHCP reservations ({len(inv['reservations'])}) ===")
    for r in inv["reservations"]:
        print(f"  {r['ip']:16} {r['mac']:18} {(r['network_name'] or '?')[:14]:14} {r['name'][:30]}")

    if inv["port_forwards"]:
        print(f"\n=== port forwards ({len(inv['port_forwards'])}) ===")
        for p in inv["port_forwards"]:
            state = "" if p["enabled"] else "  [DISABLED]"
            print(f"  {p['proto']:6} {str(p['src']):16}:{str(p['dst_port']):11} -> "
                  f"{p['fwd']}:{p['fwd_port']}   {p['name']}{state}")

    if inv["firewall_rules"]:
        print(f"\n=== firewall rules ({len(inv['firewall_rules'])}) ===")
        for r in inv["firewall_rules"]:
            state = "" if r["enabled"] else "  [DISABLED]"
            print(f"  {str(r['ruleset']):22} {str(r['action']):8} {r['name']}{state}")

    if inv["warnings"]:
        print(f"\n=== {len(inv['warnings'])} things to settle before cutover ===")
        for w in inv["warnings"]:
            n = w.get("count")
            print(f"  [{w['kind']}]" + (f"  x{n}" if n else ""))
            det = w["detail"]
            for line in (det if isinstance(det, list) else [det])[:6]:
                print(f"      {line}")
            if isinstance(det, list) and len(det) > 6:
                print(f"      ... and {len(det) - 6} more (see inventory.json)")

    n_dyn = len(inv["dynamic_clients"])
    if n_dyn:
        print(f"\n=== {n_dyn} active clients WITHOUT a reservation ===")
        print("  These hold their address only via the current lease database. A DHCP")
        print("  server swap hands them a different one — fine for phones, not fine for")
        print("  anything another host reaches by IP. Review before cutover:")
        for c in inv["dynamic_clients"][:40]:
            print(f"  {c['ip']:16} {c['mac']:18} {(c['network'] or '')[:14]:14} {c['name'][:30]}")
        if n_dyn > 40:
            print(f"  ... and {n_dyn - 40} more (see inventory.json)")


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