From 870c5cc363479ba07d03d6c9e88ccdf85c86619f Mon Sep 17 00:00:00 2001 From: ducoterra Date: Sun, 2 Aug 2026 11:09:54 -0400 Subject: [PATCH] add prefix to rules for easy identification --- AGENTS.md | 5 ++++- config.py | 3 +++ main.py | 48 +++++++++++++++++++++++++++++----------------- tests/test_main.py | 2 +- 4 files changed, 38 insertions(+), 20 deletions(-) create mode 100644 config.py diff --git a/AGENTS.md b/AGENTS.md index 3efefe2..333345e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -3,7 +3,7 @@ ## Dev commands (always run all three) ```bash -uv run pytest -v # 35 tests, all mocked (no live API calls) +uv run pytest -v # 50 tests, all mocked (no live API calls) uv run ruff check # lint uv run pyright # strict type checking ``` @@ -13,12 +13,15 @@ Order matters: fix lint/typecheck errors before touching tests. ## Architecture - `main.py` — entry point: load config → fetch public IPs → process rules → NTFY notify +- `config.py` — app constants (e.g., `RULE_NAME_PREFIX`) - `unifi_firewall.py` — UniFi Network API v1 client (TypedDicts match `network_v10.4.57_openapi.json`) - `ip_lookup.py` — fetches public IPv4/IPv6 via `curl ifconfig.me` - `config/rules.yaml` — rule definitions (zones, IPs, ports, actions) API base: `{UNIFI_HOST}/proxy/network/integration/v1/sites/{UNIFI_SITE_ID}`. Auth: `X-API-Key` header. +Rule names from `config/rules.yaml` are prefixed with `AUTO-` at runtime (`config.py:RULE_NAME_PREFIX`). First run migrates any existing unprefixed rules by renaming them. + ## Env vars Required: `UNIFI_HOST`, `UNIFI_SITE_ID`, `UNIFI_API_TOKEN` diff --git a/config.py b/config.py new file mode 100644 index 0000000..a017297 --- /dev/null +++ b/config.py @@ -0,0 +1,3 @@ +"""Application configuration constants.""" + +RULE_NAME_PREFIX = "AUTO-" diff --git a/main.py b/main.py index 1b2bacb..8d784f6 100644 --- a/main.py +++ b/main.py @@ -28,6 +28,7 @@ import requests import yaml from dotenv import load_dotenv +from config import RULE_NAME_PREFIX from ip_lookup import get_ipv4, get_ipv6 from unifi_firewall import ( FirewallPolicy, @@ -260,11 +261,12 @@ def process_rule( ) -> RuleChange: """Process a single firewall rule: create, update, or skip.""" rule_name = rule.get("name", "unnamed") + prefixed_name = f"{RULE_NAME_PREFIX}{rule_name}" 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) + logger.info("=== Processing rule: %s ===", prefixed_name) # Get zone IDs src_zone_id = zone_map.get(source_zone) @@ -273,36 +275,46 @@ def process_rule( 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} + return {"rule_name": prefixed_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} + return {"rule_name": prefixed_name, "action": "failed", "ip": public_ip, "error": error_msg} - # Find existing policy by name + # Find existing policy by prefixed name try: - existing_policies = list_policies(session, host, site_id, name_filter=rule_name) + existing_policies = list_policies(session, host, site_id, name_filter=prefixed_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} + return {"rule_name": prefixed_name, "action": "failed", "ip": public_ip, "error": error_msg} existing_policy = existing_policies[0] if existing_policies else None + # Migration: if not found by prefixed name, search for unprefixed name + if not existing_policy: + try: + existing_policies = list_policies(session, host, site_id, name_filter=rule_name) + existing_policy = existing_policies[0] if existing_policies else None + except Exception as e: # noqa: BLE001 + error_msg = f"Failed to list policies: {e}" + logger.error(error_msg) + return {"rule_name": prefixed_name, "action": "failed", "ip": public_ip, "error": error_msg} + # 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} + logger.info("Rule '%s' already matches config, skipping", prefixed_name) + return {"rule_name": prefixed_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) + prefixed_name, existing_ip, public_ip) try: payload = build_policy_payload( - name=rule_name, + name=prefixed_name, source_ip=public_ip, ip_version=ip_version, # type: ignore[arg-type] src_zone_id=src_zone_id, @@ -316,18 +328,18 @@ def process_rule( 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} + logger.info("Rule '%s' updated successfully", prefixed_name) + return {"rule_name": prefixed_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} + return {"rule_name": prefixed_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) + logger.info("Rule '%s' does not exist, creating with IP %s", prefixed_name, public_ip) try: payload = build_policy_payload( - name=rule_name, + name=prefixed_name, source_ip=public_ip, ip_version=ip_version, # type: ignore[arg-type] src_zone_id=src_zone_id, @@ -341,12 +353,12 @@ def process_rule( 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} + logger.info("Rule '%s' created successfully", prefixed_name) + return {"rule_name": prefixed_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} + return {"rule_name": prefixed_name, "action": "failed", "ip": public_ip, "error": error_msg} def main() -> None: diff --git a/tests/test_main.py b/tests/test_main.py index 2b285b2..0ed2449 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -126,7 +126,7 @@ class TestProcessRule: change = process_rule(mock_session, "https://unifi.local", "site-1", rule, zone_map, "1.2.3.4") assert change["action"] == "created" - assert change["rule_name"] == "Test Rule" + assert change["rule_name"] == "AUTO-Test Rule" assert change["error"] is None def test_skips_rule_with_matching_config(self) -> None: