#!/usr/bin/env python3
"""Create a zone-based firewall policy on one UniFi console that allows TCP from a
source zone to one destination IP on given ports (e.g. guest WiFi -> a LAN web host).

Usage:
    docs/examples/waters-church-na/zone-allow-policy.py north-attleboro Hotspot 10.1.1.4 80,443 \
        --name "Guest to translate.watersmedia.org (Boba-Fett web)"            # dry run
    docs/examples/waters-church-na/zone-allow-policy.py north-attleboro Hotspot 10.1.1.4 80,443 \
        --name "Guest to translate.watersmedia.org (Boba-Fett web)" --apply    # create it
    docs/examples/waters-church-na/zone-allow-policy.py north-attleboro --list-zones

Without --apply nothing is written: the script prints the zones it resolved, any
existing policy with the same name, and the exact JSON it would POST. The
destination zone is found from the console's networks (the one whose subnet
contains the IP). User-defined policies are evaluated before the system ones
(index >= 30000), so the allow lands above "Post-Authorization Restrictions".

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 ipaddress
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


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]:
        data = self.request("GET", f"/integration/v1/sites/{sid}/firewall/zones?limit=200")["data"]
        return {z["name"]: z["id"] for z in data}

    def zone_for_ip(self, sid: str, ip: str) -> tuple[str, str]:
        """Zone id and network name whose subnet contains ip (via the legacy networkconf)."""
        nets = {n["id"]: n for n in self.request("GET", f"/integration/v1/sites/{sid}/networks?limit=200")["data"]}
        addr = ipaddress.ip_address(ip)
        for n in self.request("GET", "/api/s/default/rest/networkconf")["data"]:
            subnet = n.get("ip_subnet")
            if not subnet:
                continue
            if addr in ipaddress.ip_interface(subnet).network:
                zone_id = nets.get(n.get("external_id", ""), {}).get("zoneId")
                if zone_id:
                    return zone_id, n["name"]
        raise SystemExit(f"no gateway network on {self.name} contains {ip}")

    def policies(self, sid: str) -> list[dict]:
        return self.request("GET", f"/integration/v1/sites/{sid}/firewall/policies?limit=500")["data"]


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("site", choices=sorted(SITES))
    ap.add_argument("source_zone", nargs="?", help="source zone name, e.g. Hotspot")
    ap.add_argument("dest_ip", nargs="?", help="destination IPv4 address")
    ap.add_argument("ports", nargs="?", help="comma-separated TCP ports, e.g. 80,443")
    ap.add_argument("--name", help="policy name (default is generated)")
    ap.add_argument("--apply", action="store_true", help="actually create the policy")
    ap.add_argument("--list-zones", action="store_true", help="print the console's zones and exit")
    a = ap.parse_args()

    c = Console(a.site)
    sid = c.site_id()
    zones = c.zones(sid)
    if a.list_zones:
        for name, zid in sorted(zones.items()):
            print(f"{name:12} {zid}")
        return 0
    if not (a.source_zone and a.dest_ip and a.ports):
        ap.error("source_zone, dest_ip and ports are required unless --list-zones")
    if a.source_zone not in zones:
        raise SystemExit(f"zone {a.source_zone!r} not on {a.site}; have: {', '.join(sorted(zones))}")
    ports = [int(p) for p in a.ports.split(",")]
    dest_zone, dest_net = c.zone_for_ip(sid, a.dest_ip)
    dest_zone_name = next(n for n, z in zones.items() if z == dest_zone)
    name = a.name or f"{a.source_zone} to {a.dest_ip} tcp {a.ports}"

    body = {
        "name": name,
        "enabled": True,
        "action": {"type": "ALLOW", "allowReturnTraffic": True},
        "source": {"zoneId": zones[a.source_zone]},
        "destination": {
            "zoneId": dest_zone,
            "trafficFilter": {
                "type": "IP_ADDRESS",
                "ipAddressFilter": {"type": "IP_ADDRESSES", "matchOpposite": False,
                                    "items": [{"type": "IP_ADDRESS", "value": a.dest_ip}]},
                "portFilter": {"type": "PORTS", "matchOpposite": False,
                               "items": [{"type": "PORT_NUMBER", "value": p} for p in ports]},
            },
        },
        "ipProtocolScope": {"ipVersion": "IPV4",
                            "protocolFilter": {"type": "NAMED_PROTOCOL", "protocol": {"name": "TCP"},
                                               "matchOpposite": False}},
        "loggingEnabled": False,
    }

    print(f"{a.site} ({c.gateway}): {a.source_zone} -> {dest_zone_name} ({dest_net}) {a.dest_ip} tcp {a.ports}")
    existing = [p for p in c.policies(sid) if p.get("name") == name]
    if existing:
        print(f"a policy named {name!r} already exists (id {existing[0]['id']}, index {existing[0].get('index')}); nothing to do")
        return 0
    if not a.apply:
        print("DRY RUN, would POST /integration/v1/sites/<site>/firewall/policies:")
        print(json.dumps(body, indent=2))
        return 0
    created = c.request("POST", f"/integration/v1/sites/{sid}/firewall/policies", body)
    print(f"CREATED id {created.get('id')} index {created.get('index')} enabled {created.get('enabled')}")
    blockers = [p for p in c.policies(sid) if p.get("name") == "Post-Authorization Restrictions"
                and p["source"]["zoneId"] == zones[a.source_zone] and p["destination"]["zoneId"] == dest_zone]
    for b in blockers:
        order = "above" if created.get("index", 0) < b.get("index", 0) else "BELOW"
        print(f"new policy sits {order} 'Post-Authorization Restrictions' (index {b.get('index')})")
    return 0


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