init

main.py

@@ -0,0 +1,371 @@
+"""
+UniFi firewall rule updater.
+
+Updates UniFi Dream Machine firewall rules with the host machine's public
+IPv4 and IPv6 addresses via the UniFi Network API.
+
+Exported env vars:
+    UNIFI_HOST - UniFi controller URL (e.g., https://10.1.0.1)
+    UNIFI_SITE_ID - UniFi site ID
+    UNIFI_API_TOKEN - UniFi API token
+    UNIFI_VERIFY_SSL - verify SSL certificates (default: false)
+    CONFIG_FILE - path to YAML file with firewall rules (default: config/rules.yaml)
+    LOG_LEVEL - logging level (DEBUG, INFO, WARNING, ERROR; default: INFO)
+    DEBUG - if true, starts debugpy and waits for a remote debugger connection on port 5678
+    NTFY_URL - NTFY notification server URL
+    NTFY_TOPIC - NTFY topic to send notifications to
+    NTFY_API_KEY - NTFY API key for authentication
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+import logging
+from typing import Literal, TypedDict
+
+import requests
+import yaml
+from dotenv import load_dotenv
+
+from ip_lookup import get_ipv4, get_ipv6
+from unifi_firewall import (
+    FirewallZone,
+    build_policy_payload,
+    create_policy,
+    get_session,
+    get_source_ip_from_policy,
+    list_policies,
+    list_zones,
+    update_policy,
+)
+
+load_dotenv()
+
+DEBUG = os.getenv("DEBUG", "false").lower() == "true"
+if DEBUG:
+    import debugpy  # noqa: T100
+
+    debugpy.listen(("0.0.0.0", 5678))  # noqa: T100
+    print("DEBUG: debugpy listening on 0.0.0.0:5678, waiting for debugger to attach...")
+    debugpy.wait_for_client()  # noqa: T100
+    print("DEBUG: debugger attached, resuming execution...")
+
+
+try:
+    from yaml import CLoader as Loader
+except ImportError:
+    from yaml import Loader  # type: ignore[attr-defined]
+
+log_level = os.getenv("LOG_LEVEL", "INFO").upper()
+
+logging.basicConfig(
+    level=log_level,
+    format="%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger = logging.getLogger(__name__)
+logger.setLevel(log_level)
+
+UNIFI_HOST = os.getenv("UNIFI_HOST")
+UNIFI_SITE_ID = os.getenv("UNIFI_SITE_ID")
+UNIFI_API_TOKEN = os.getenv("UNIFI_API_TOKEN")
+UNIFI_VERIFY_SSL = os.getenv("UNIFI_VERIFY_SSL", "false").lower() == "true"
+
+CONFIG_FILE = os.getenv("CONFIG_FILE", "config/rules.yaml")
+
+NTFY_URL = os.getenv("NTFY_URL", "")
+NTFY_TOPIC = os.getenv("NTFY_TOPIC", "")
+NTFY_API_KEY = os.getenv("NTFY_API_KEY", "")
+
+
+class PortRange(TypedDict):
+    """Port range definition."""
+
+    start: int
+    stop: int
+
+
+class RuleConfig(TypedDict, total=False):
+    """Firewall rule configuration from YAML."""
+
+    name: str
+    source_zone: str
+    dest_zone: str
+    ip_version: Literal["IPV4", "IPV6"]
+    action: Literal["ALLOW", "BLOCK", "REJECT"]
+    allow_return_traffic: bool
+    protocol: str | None
+    dest_ports: list[int] | None
+    dest_port_ranges: list[PortRange] | None
+    logging_enabled: bool
+    enabled: bool
+
+
+class RulesConfig(TypedDict):
+    """Top-level rules configuration."""
+
+    rules: list[RuleConfig]
+
+
+class RuleChange(TypedDict):
+    """Record of a rule change."""
+
+    rule_name: str
+    action: Literal["created", "updated", "skipped", "failed"]
+    ip: str
+    error: str | None
+
+
+def send_ntfy_notification(title: str, message: str, priority: int = 3) -> None:
+    """Send an NTFY notification."""
+    if not NTFY_URL or not NTFY_TOPIC:
+        return
+    try:
+        headers = {
+            "Title": title,
+            "Priority": str(priority),
+        }
+        if NTFY_API_KEY:
+            headers["Authorization"] = f"Bearer {NTFY_API_KEY}"
+        logger.info("Sending NTFY notification: %s", title)
+        response = requests.post(
+            f"{NTFY_URL}/{NTFY_TOPIC}",
+            data=message.encode(),
+            headers=headers,
+            timeout=10,
+        )
+        response.raise_for_status()
+        logger.info("NTFY notification sent: %s", title)
+    except requests.RequestException as e:
+        logger.warning("Failed to send NTFY notification: %s", e)
+
+
+def load_config(path: str) -> list[RuleConfig]:
+    """Load firewall rules from YAML config file."""
+    logger.info("Loading config file: %s", path)
+    try:
+        with open(path) as f:
+            config: RulesConfig = yaml.load(f, Loader)
+        rules = config.get("rules", [])
+        logger.debug("Loaded %d rule(s) from config", len(rules))
+        return rules
+    except FileNotFoundError as e:
+        logger.error("Config file not found: %s", e)
+        sys.exit(1)
+    except yaml.YAMLError as e:
+        logger.error("Failed to parse config file: %s", e)
+        sys.exit(1)
+
+
+def build_zone_map(zones: list[FirewallZone]) -> dict[str, str]:
+    """Build a mapping of zone names to zone IDs."""
+    zone_map: dict[str, str] = {}
+    for zone in zones:
+        zone_map[zone["name"]] = zone["id"]
+    logger.debug("Zone map: %s", zone_map)
+    return zone_map
+
+
+def process_rule(
+    session: requests.Session,
+    host: str,
+    site_id: str,
+    rule: RuleConfig,
+    zone_map: dict[str, str],
+    public_ip: str,
+) -> RuleChange:
+    """Process a single firewall rule: create, update, or skip."""
+    rule_name = rule.get("name", "unnamed")
+    ip_version = rule.get("ip_version", "IPV4")
+    source_zone = rule.get("source_zone", "WAN")
+    dest_zone = rule.get("dest_zone", "LAN")
+
+    logger.info("=== Processing rule: %s ===", rule_name)
+
+    # Get zone IDs
+    src_zone_id = zone_map.get(source_zone)
+    dst_zone_id = zone_map.get(dest_zone)
+
+    if not src_zone_id:
+        error_msg = f"Source zone '{source_zone}' not found"
+        logger.error(error_msg)
+        return {"rule_name": rule_name, "action": "failed", "ip": public_ip, "error": error_msg}
+
+    if not dst_zone_id:
+        error_msg = f"Destination zone '{dest_zone}' not found"
+        logger.error(error_msg)
+        return {"rule_name": rule_name, "action": "failed", "ip": public_ip, "error": error_msg}
+
+    # Find existing policy by name
+    try:
+        existing_policies = list_policies(session, host, site_id, name_filter=rule_name)
+    except Exception as e:  # noqa: BLE001
+        error_msg = f"Failed to list policies: {e}"
+        logger.error(error_msg)
+        return {"rule_name": rule_name, "action": "failed", "ip": public_ip, "error": error_msg}
+
+    existing_policy = existing_policies[0] if existing_policies else None
+
+    # Check if policy exists and IP matches
+    if existing_policy:
+        existing_ip = get_source_ip_from_policy(existing_policy)
+        if existing_ip == public_ip:
+            logger.info("Rule '%s' already has correct IP (%s), skipping", rule_name, public_ip)
+            return {"rule_name": rule_name, "action": "skipped", "ip": public_ip, "error": None}
+
+        # Update existing policy
+        logger.info("Rule '%s' exists with IP %s, updating to %s", rule_name, existing_ip, public_ip)
+        try:
+            payload = build_policy_payload(
+                name=rule_name,
+                source_ip=public_ip,
+                ip_version=ip_version,  # type: ignore[arg-type]
+                src_zone_id=src_zone_id,
+                dst_zone_id=dst_zone_id,
+                action_type=rule.get("action", "ALLOW"),  # type: ignore[arg-type]
+                allow_return_traffic=rule.get("allow_return_traffic", True),
+                protocol=rule.get("protocol"),
+                dest_ports=rule.get("dest_ports"),
+                dest_port_ranges=rule.get("dest_port_ranges"),  # type: ignore[arg-type]
+                logging_enabled=rule.get("logging_enabled", False),
+                enabled=rule.get("enabled", True),
+            )
+            update_policy(session, host, site_id, existing_policy["id"], payload)
+            logger.info("Rule '%s' updated successfully", rule_name)
+            return {"rule_name": rule_name, "action": "updated", "ip": public_ip, "error": None}
+        except Exception as e:  # noqa: BLE001
+            error_msg = f"Failed to update rule: {e}"
+            logger.error(error_msg)
+            return {"rule_name": rule_name, "action": "failed", "ip": public_ip, "error": error_msg}
+
+    # Create new policy
+    logger.info("Rule '%s' does not exist, creating with IP %s", rule_name, public_ip)
+    try:
+        payload = build_policy_payload(
+            name=rule_name,
+            source_ip=public_ip,
+            ip_version=ip_version,  # type: ignore[arg-type]
+            src_zone_id=src_zone_id,
+            dst_zone_id=dst_zone_id,
+            action_type=rule.get("action", "ALLOW"),  # type: ignore[arg-type]
+            allow_return_traffic=rule.get("allow_return_traffic", True),
+            protocol=rule.get("protocol"),
+            dest_ports=rule.get("dest_ports"),
+            dest_port_ranges=rule.get("dest_port_ranges"),  # type: ignore[arg-type]
+            logging_enabled=rule.get("logging_enabled", False),
+            enabled=rule.get("enabled", True),
+        )
+        create_policy(session, host, site_id, payload)
+        logger.info("Rule '%s' created successfully", rule_name)
+        return {"rule_name": rule_name, "action": "created", "ip": public_ip, "error": None}
+    except Exception as e:  # noqa: BLE001
+        error_msg = f"Failed to create rule: {e}"
+        logger.error(error_msg)
+        return {"rule_name": rule_name, "action": "failed", "ip": public_ip, "error": error_msg}
+
+
+def main() -> None:
+    """Main entry point."""
+    logger.info("=== UniFi Firewall Update Starting ===")
+    logger.debug("Log level: %s", log_level)
+    logger.debug("UNIFI_HOST: %s", UNIFI_HOST)
+    logger.debug("UNIFI_SITE_ID: %s", UNIFI_SITE_ID)
+    logger.debug("UNIFI_API_TOKEN: %s", "****" if UNIFI_API_TOKEN else "not set")
+    logger.debug("UNIFI_VERIFY_SSL: %s", UNIFI_VERIFY_SSL)
+    logger.debug("CONFIG_FILE: %s", CONFIG_FILE)
+
+    # Validate required env vars
+    if not all([UNIFI_HOST, UNIFI_SITE_ID, UNIFI_API_TOKEN]):
+        logger.error("UNIFI_HOST, UNIFI_SITE_ID, and UNIFI_API_TOKEN must be set!")
+        sys.exit(1)
+
+    assert UNIFI_HOST is not None
+    assert UNIFI_SITE_ID is not None
+    assert UNIFI_API_TOKEN is not None
+
+    # Load config
+    rules = load_config(CONFIG_FILE)
+    if not rules:
+        logger.warning("No rules found in config file")
+        return
+
+    # Fetch public IPs
+    public_ipv4 = get_ipv4()
+    if not public_ipv4:
+        logger.error("Failed to fetch public IPv4 address")
+        sys.exit(1)
+    logger.info("Public IPv4: %s", public_ipv4)
+
+    public_ipv6 = get_ipv6()
+    if not public_ipv6:
+        logger.error("Failed to fetch public IPv6 address")
+        sys.exit(1)
+    logger.info("Public IPv6: %s", public_ipv6)
+
+    # Create session and list zones
+    session = get_session(UNIFI_HOST, UNIFI_API_TOKEN, UNIFI_VERIFY_SSL)
+    try:
+        zones = list_zones(session, UNIFI_HOST, UNIFI_SITE_ID)
+    except Exception as e:  # noqa: BLE001
+        logger.error("Failed to list firewall zones: %s", e)
+        sys.exit(1)
+
+    zone_map = build_zone_map(zones)
+
+    # Process rules
+    changes: list[RuleChange] = []
+
+    for rule in rules:
+        ip_version = rule.get("ip_version", "IPV4")
+        public_ip = public_ipv4 if ip_version == "IPV4" else public_ipv6
+
+        if not public_ip:
+            logger.warning("Skipping rule '%s': no %s address available", rule.get("name"), ip_version)
+            continue
+
+        change = process_rule(session, UNIFI_HOST, UNIFI_SITE_ID, rule, zone_map, public_ip)
+        changes.append(change)
+
+    # Summary
+    created = [c for c in changes if c["action"] == "created"]
+    updated = [c for c in changes if c["action"] == "updated"]
+    skipped = [c for c in changes if c["action"] == "skipped"]
+    failed = [c for c in changes if c["action"] == "failed"]
+
+    logger.info("=== Summary ===")
+    logger.info("Created: %d", len(created))
+    logger.info("Updated: %d", len(updated))
+    logger.info("Skipped: %d", len(skipped))
+    logger.info("Failed: %d", len(failed))
+
+    # Build NTFY message
+    ntfy_lines: list[str] = []
+    for c in created:
+        ntfy_lines.append(f"+ {c['rule_name']}: {c['ip']}")
+    for c in updated:
+        ntfy_lines.append(f"~ {c['rule_name']}: {c['ip']}")
+    for c in failed:
+        ntfy_lines.append(f"! {c['rule_name']}: {c['error']}")
+
+    if ntfy_lines:
+        ntfy_message = "\n".join(ntfy_lines)
+        has_changes = bool(created or updated)
+        has_failures = bool(failed)
+        if has_failures:
+            ntfy_title = "Firewall Update Failed"
+            ntfy_priority = 4
+        elif has_changes:
+            ntfy_title = "Firewall Rules Updated"
+            ntfy_priority = 4
+        else:
+            ntfy_title = "Firewall Rules Unchanged"
+            ntfy_priority = 2
+        send_ntfy_notification(ntfy_title, ntfy_message, priority=ntfy_priority)
+
+    logger.info("=== UniFi Firewall Update Complete ===")
+
+
+if __name__ == "__main__":
+    main()