services/firewall
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
e360546de9f87c6401f136f10ac53c198c354c40
firewall/main.py
T
ducoterra e360546de9
Build and Push Container / build-and-push (push) Successful in 19s
Details
init
2026-08-01 19:07:05 -04:00

372 lines
12 KiB
Python
Raw Blame History

"""
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()
Reference in New Issue View Git Blame Copy Permalink