firewall/main.py

"""
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 logging
import os
import sys
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 (
    FirewallPolicy,
    FirewallZone,
    build_policy_payload,
    create_policy,
    get_action_from_policy,
    get_dest_ports_from_policy,
    get_protocol_from_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
    description: str | None


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 policy_matches_config(
    existing_policy: FirewallPolicy,
    rule: RuleConfig,
    zone_map: dict[str, str],
    public_ip: str,
) -> bool:
    """Check if an existing policy matches the desired config exactly."""
    # 1. Source IP
    existing_ip = get_source_ip_from_policy(existing_policy)
    if existing_ip != public_ip:
        return False

    # 2. Dest ports (compare as sorted lists)
    existing_ports_result = get_dest_ports_from_policy(existing_policy)
    rule_dest_ports = rule.get("dest_ports") or []
    if existing_ports_result is None:
        if rule_dest_ports:
            return False
    else:
        existing_ports, existing_ranges = existing_ports_result
        if sorted(existing_ports) != sorted(rule_dest_ports):
            return False

        # 3. Dest port ranges
        rule_dest_ranges = rule.get("dest_port_ranges") or []
        if len(existing_ranges) != len(rule_dest_ranges):
            return False
        for er, rr in zip(sorted(existing_ranges, key=lambda x: x["start"]), sorted(rule_dest_ranges, key=lambda x: x["start"])):
            if er["start"] != rr["start"] or er["stop"] != rr["stop"]:
                return False

    # 4. Protocol (case-insensitive)
    existing_protocol = get_protocol_from_policy(existing_policy)
    rule_protocol = rule.get("protocol")
    if existing_protocol is not None and rule_protocol is not None:
        if existing_protocol.upper() != rule_protocol.upper():
            return False
    elif existing_protocol is not None or rule_protocol is not None:
        return False

    # 5. Action type
    existing_action_result = get_action_from_policy(existing_policy)
    if existing_action_result is None:
        return False
    existing_action_type, existing_allow_return = existing_action_result
    if existing_action_type != rule.get("action", "ALLOW"):
        return False

    # 6. allowReturnTraffic
    if existing_action_type == "ALLOW" and existing_allow_return != rule.get("allow_return_traffic", True):
        return False

    # 7. Source zone ID
    rule_source_zone = rule.get("source_zone", "WAN")
    if existing_policy["source"]["zoneId"] != zone_map.get(rule_source_zone):
        return False

    # 8. Dest zone ID
    rule_dest_zone = rule.get("dest_zone", "LAN")
    if existing_policy["destination"]["zoneId"] != zone_map.get(rule_dest_zone):
        return False

    # 9. IP version
    if existing_policy["ipProtocolScope"]["ipVersion"] != rule.get("ip_version", "IPV4"):
        return False

    # 10. enabled
    if existing_policy["enabled"] != rule.get("enabled", True):
        return False

    # 11. loggingEnabled
    if existing_policy["loggingEnabled"] != rule.get("logging_enabled", False):
        return False

    # 12. description
    return existing_policy["description"] == rule.get("description")


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 matches config
    if existing_policy:
        if policy_matches_config(existing_policy, rule, zone_map, public_ip):
            logger.info("Rule '%s' already matches config, skipping", rule_name)
            return {"rule_name": rule_name, "action": "skipped", "ip": public_ip, "error": None}

        # Update existing policy (IP or other attributes changed)
        existing_ip = get_source_ip_from_policy(existing_policy)
        logger.info("Rule '%s' exists but config changed (current IP: %s, desired IP: %s), updating",
                    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()