#!/usr/bin/env python3
"""Remove a subnet from the guest "Post-Authorization Restrictions" on one UniFi console.

UniFi blocks guest (Hotspot zone) clients from the private ranges listed in the
guest hotspot's Post-Authorization Restrictions. That list is enforced in two
places: as system firewall policies on the gateway (one per zone pair, which a
user policy can sit above) and on the access points themselves, which a gateway
policy cannot override. To let guests reach one internal host you therefore
remove its range from the restriction list and rely on the zone firewall's
default Hotspot -> Internal block plus an explicit allow (zone-allow-policy.py).

Usage:
    docs/examples/waters-church-na/guest-restrictions.py north-attleboro --show
    docs/examples/waters-church-na/guest-restrictions.py north-attleboro --remove 10.0.0.0/8            # dry run
    docs/examples/waters-church-na/guest-restrictions.py north-attleboro --remove 10.0.0.0/8 --apply

It updates every "Post-Authorization Restrictions" policy whose source is the
Hotspot zone through the Network Integration API, then the legacy
guest_access setting (restricted_subnet_1..3) that the access points read.

Needs the per-console Network API key at ~/.config/unifi/<site>-network-api-key
(see unifi-dns-record.py). The key is only ever sent in the X-API-KEY header.
"""

from __future__ import annotations

import argparse
import json
import os
import ssl
import sys
import urllib.error
import urllib.request

SITES = {
    "apollo-beach": "10.3.0.1",
    "north-attleboro": "10.1.0.1",
    "woonsocket": "10.2.0.1",
    "norwood": "10.4.0.1",
    "southcoast": "10.5.0.1",
}
KEY_DIR = os.path.expanduser("~/.config/unifi")
TIMEOUT = 20
POLICY_NAME = "Post-Authorization Restrictions"


class Console:
    def __init__(self, name: str):
        self.name = name
        self.gateway = SITES[name]
        with open(os.path.join(KEY_DIR, f"{name}-network-api-key")) as f:
            self.key = f.read().strip()
        self.ctx = ssl.create_default_context()
        self.ctx.check_hostname = False
        self.ctx.verify_mode = ssl.CERT_NONE

    def request(self, method: str, path: str, body: dict | None = None):
        url = f"https://{self.gateway}/proxy/network{path}"
        data = json.dumps(body).encode() if body is not None else None
        req = urllib.request.Request(url, data=data, method=method, headers={
            "X-API-KEY": self.key, "Accept": "application/json",
            "Content-Type": "application/json"})
        try:
            with urllib.request.urlopen(req, timeout=TIMEOUT, context=self.ctx) as r:
                raw = r.read().decode()
        except urllib.error.HTTPError as e:
            raise SystemExit(f"{method} {path} -> HTTP {e.code}: {e.read().decode(errors='replace')[:600]}")
        return json.loads(raw) if raw else None

    def site_id(self) -> str:
        return self.request("GET", "/integration/v1/sites")["data"][0]["id"]

    def zones(self, sid: str) -> dict[str, str]:
        return {z["id"]: z["name"] for z in self.request("GET", f"/integration/v1/sites/{sid}/firewall/zones?limit=200")["data"]}

    def restriction_policies(self, sid: str, zones: dict[str, str]) -> list[dict]:
        pols = self.request("GET", f"/integration/v1/sites/{sid}/firewall/policies?limit=500")["data"]
        return [p for p in pols if p.get("name") == POLICY_NAME and zones.get(p["source"]["zoneId"]) == "Hotspot"]

    def guest_access(self) -> dict:
        for s in self.request("GET", "/api/s/default/get/setting")["data"]:
            if s.get("key") == "guest_access":
                return s
        raise SystemExit("no guest_access setting found")


def items_of(p: dict) -> list[str]:
    f = (p.get("destination") or {}).get("trafficFilter") or {}
    return [i["value"] for i in (f.get("ipAddressFilter") or {}).get("items", [])]


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("site", choices=sorted(SITES))
    ap.add_argument("--show", action="store_true", help="print the current restrictions and exit")
    ap.add_argument("--remove", help="subnet to remove, e.g. 10.0.0.0/8")
    ap.add_argument("--apply", action="store_true", help="actually write the change")
    a = ap.parse_args()

    c = Console(a.site)
    sid = c.site_id()
    zones = c.zones(sid)
    pols = c.restriction_policies(sid, zones)
    ga = c.guest_access()
    legacy = {k: v for k, v in ga.items() if k.startswith("restricted_subnet_")}

    print(f"{a.site} ({c.gateway}) guest Post-Authorization Restrictions:")
    for p in pols:
        print(f"  gateway policy Hotspot -> {zones.get(p['destination']['zoneId']):8} index {p['index']}: {', '.join(items_of(p))}")
    print(f"  legacy guest_access (access points): {json.dumps(legacy)}")
    if a.show or not a.remove:
        return 0

    targets = [p for p in pols if a.remove in items_of(p)]
    legacy_hits = {k: v for k, v in legacy.items() if v == a.remove}
    if not targets and not legacy_hits:
        print(f"\n{a.remove} is not in any restriction list; nothing to do")
        return 0
    print(f"\nwould remove {a.remove} from legacy field(s) {list(legacy_hits)}; the {len(targets)} gateway policy(ies) listing it regenerate from that setting")
    if not a.apply:
        print("DRY RUN, add --apply to write it")
        return 0

    # The gateway policies are SYSTEM_DEFINED and the Integration API refuses to
    # edit them (api.firewall.policy.operation.system-defined-firewall-policy-update-forbidden).
    # They are generated from the legacy guest_access restricted_subnet_* fields, so
    # that setting is the one lever; the policies and the APs follow it.
    if not legacy_hits:
        print("  no legacy restricted_subnet field holds it; the gateway policies cannot be edited directly")
        return 1
    # Re-pack the remaining subnets into slots 1..n and blank the rest.
    remaining = [v for k, v in sorted(legacy.items()) if v and v != a.remove]
    payload = {f"restricted_subnet_{i}": (remaining[i - 1] if i <= len(remaining) else "") for i in range(1, 4)}
    print(f"  POST /api/s/default/set/setting/guest_access {json.dumps(payload)}")
    c.request("POST", "/api/s/default/set/setting/guest_access", payload)
    ga2 = c.guest_access()
    print("  legacy guest_access now:", json.dumps({k: v for k, v in ga2.items() if k.startswith("restricted_subnet_")}))
    for p in c.restriction_policies(sid, zones):
        print(f"  gateway policy Hotspot -> {zones.get(p['destination']['zoneId']):8} now: {', '.join(items_of(p))}")
    print("done; access points re-provision within a minute, guests should reconnect")
    return 0


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