From a81dd4f3df5d64ad3446e8481129593bbf2fe933 Mon Sep 17 00:00:00 2001 From: ducoterra Date: Sat, 1 Aug 2026 19:50:32 -0400 Subject: [PATCH] fix issue where you couldn't update ports --- main.py | 99 +++++++++++- tests/test_main.py | 287 ++++++++++++++++++++++++++++++++++- tests/test_unifi_firewall.py | 126 +++++++++++++++ unifi_firewall.py | 84 ++++++++++ 4 files changed, 587 insertions(+), 9 deletions(-) diff --git a/main.py b/main.py index 225f221..355681a 100644 --- a/main.py +++ b/main.py @@ -19,10 +19,9 @@ Exported env vars: from __future__ import annotations +import logging import os import sys - -import logging from typing import Literal, TypedDict import requests @@ -31,9 +30,13 @@ 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, @@ -101,6 +104,7 @@ class RuleConfig(TypedDict, total=False): dest_port_ranges: list[PortRange] | None logging_enabled: bool enabled: bool + description: str | None class RulesConfig(TypedDict): @@ -168,6 +172,84 @@ def build_zone_map(zones: list[FirewallZone]) -> dict[str, str]: 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, @@ -208,15 +290,16 @@ def process_rule( existing_policy = existing_policies[0] if existing_policies else None - # Check if policy exists and IP matches + # Check if policy exists and matches config 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) + 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 - logger.info("Rule '%s' exists with IP %s, updating to %s", rule_name, existing_ip, public_ip) + # 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, diff --git a/tests/test_main.py b/tests/test_main.py index 3248657..2b285b2 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -9,6 +9,7 @@ from main import ( RuleConfig, build_zone_map, load_config, + policy_matches_config, process_rule, send_ntfy_notification, ) @@ -128,7 +129,7 @@ class TestProcessRule: assert change["rule_name"] == "Test Rule" assert change["error"] is None - def test_skips_rule_with_matching_ip(self) -> None: + def test_skips_rule_with_matching_config(self) -> None: mock_session = MagicMock() existing_policy = { "id": "pol-1", @@ -143,6 +144,14 @@ class TestProcessRule: }, }, }, + "destination": {"zoneId": "lan-id", "trafficFilter": None}, + "action": {"type": "ALLOW", "allowReturnTraffic": True}, + "ipProtocolScope": {"ipVersion": "IPV4", "protocolFilter": None}, + "enabled": True, + "loggingEnabled": False, + "index": 0, + "metadata": {}, + "description": None, } mock_session.get.return_value = MagicMock(json=MagicMock(return_value={"data": [existing_policy]}), raise_for_status=MagicMock()) @@ -224,3 +233,279 @@ class TestProcessRule: assert change["action"] == "failed" assert "Failed to list policies" in (change["error"] or "") + + def test_updates_rule_with_different_ports_same_ip(self) -> None: + mock_session = MagicMock() + existing_policy = { + "id": "pol-1", + "name": "Test Rule", + "source": { + "zoneId": "wan-id", + "trafficFilter": { + "type": "IP_ADDRESS", + "ipAddressFilter": { + "type": "IP_ADDRESSES", + "items": [{"type": "IP_ADDRESS", "value": "1.2.3.4"}], + }, + }, + }, + "destination": { + "zoneId": "lan-id", + "trafficFilter": { + "type": "PORT", + "portFilter": { + "type": "PORTS", + "items": [ + {"type": "PORT_NUMBER", "value": 80}, + {"type": "PORT_NUMBER", "value": 443}, + ], + }, + }, + }, + "action": {"type": "ALLOW", "allowReturnTraffic": True}, + "ipProtocolScope": {"ipVersion": "IPV4", "protocolFilter": None}, + "enabled": True, + "loggingEnabled": False, + "index": 0, + "metadata": {}, + "description": None, + } + mock_session.get.return_value = MagicMock(json=MagicMock(return_value={"data": [existing_policy]}), raise_for_status=MagicMock()) + mock_session.put.return_value = MagicMock(json=MagicMock(return_value={"id": "pol-1"}), raise_for_status=MagicMock()) + + rule: RuleConfig = { + "name": "Test Rule", + "source_zone": "WAN", + "dest_zone": "LAN", + "ip_version": "IPV4", + "action": "ALLOW", + "dest_ports": [80, 443, 8080], + } + zone_map = {"WAN": "wan-id", "LAN": "lan-id"} + + change = process_rule(mock_session, "https://unifi.local", "site-1", rule, zone_map, "1.2.3.4") + + assert change["action"] == "updated" + mock_session.put.assert_called_once() + + +class TestPolicyMatchesConfig: + def test_returns_true_when_all_match(self) -> None: + existing_policy = { + "id": "pol-1", + "name": "Test Rule", + "source": { + "zoneId": "wan-id", + "trafficFilter": { + "type": "IP_ADDRESS", + "ipAddressFilter": { + "type": "IP_ADDRESSES", + "items": [{"type": "IP_ADDRESS", "value": "1.2.3.4"}], + }, + }, + }, + "destination": { + "zoneId": "lan-id", + "trafficFilter": { + "type": "PORT", + "portFilter": { + "type": "PORTS", + "items": [ + {"type": "PORT_NUMBER", "value": 80}, + {"type": "PORT_NUMBER", "value": 443}, + ], + }, + }, + }, + "action": {"type": "ALLOW", "allowReturnTraffic": True}, + "ipProtocolScope": { + "ipVersion": "IPV4", + "protocolFilter": { + "type": "NAMED_PROTOCOL", + "protocol": {"name": "TCP"}, + }, + }, + "enabled": True, + "loggingEnabled": False, + "index": 0, + "metadata": {}, + "description": "Test description", + } + rule: RuleConfig = { + "name": "Test Rule", + "source_zone": "WAN", + "dest_zone": "LAN", + "ip_version": "IPV4", + "action": "ALLOW", + "allow_return_traffic": True, + "protocol": "tcp", + "dest_ports": [443, 80], + "logging_enabled": False, + "enabled": True, + "description": "Test description", + } + zone_map = {"WAN": "wan-id", "LAN": "lan-id"} + + result = policy_matches_config(existing_policy, rule, zone_map, "1.2.3.4") # type: ignore[arg-type] + assert result is True + + def test_returns_false_different_ports(self) -> None: + existing_policy = { + "id": "pol-1", + "name": "Test Rule", + "source": { + "zoneId": "wan-id", + "trafficFilter": { + "type": "IP_ADDRESS", + "ipAddressFilter": { + "type": "IP_ADDRESSES", + "items": [{"type": "IP_ADDRESS", "value": "1.2.3.4"}], + }, + }, + }, + "destination": { + "zoneId": "lan-id", + "trafficFilter": { + "type": "PORT", + "portFilter": { + "type": "PORTS", + "items": [ + {"type": "PORT_NUMBER", "value": 80}, + {"type": "PORT_NUMBER", "value": 443}, + ], + }, + }, + }, + "action": {"type": "ALLOW", "allowReturnTraffic": True}, + "ipProtocolScope": {"ipVersion": "IPV4", "protocolFilter": None}, + "enabled": True, + "loggingEnabled": False, + "index": 0, + "metadata": {}, + "description": None, + } + rule: RuleConfig = { + "name": "Test Rule", + "source_zone": "WAN", + "dest_zone": "LAN", + "ip_version": "IPV4", + "action": "ALLOW", + "dest_ports": [80, 443, 8080], + } + zone_map = {"WAN": "wan-id", "LAN": "lan-id"} + + result = policy_matches_config(existing_policy, rule, zone_map, "1.2.3.4") # type: ignore[arg-type] + assert result is False + + def test_returns_false_different_protocol(self) -> None: + existing_policy = { + "id": "pol-1", + "name": "Test Rule", + "source": { + "zoneId": "wan-id", + "trafficFilter": { + "type": "IP_ADDRESS", + "ipAddressFilter": { + "type": "IP_ADDRESSES", + "items": [{"type": "IP_ADDRESS", "value": "1.2.3.4"}], + }, + }, + }, + "destination": {"zoneId": "lan-id", "trafficFilter": None}, + "action": {"type": "ALLOW", "allowReturnTraffic": True}, + "ipProtocolScope": { + "ipVersion": "IPV4", + "protocolFilter": { + "type": "NAMED_PROTOCOL", + "protocol": {"name": "TCP"}, + }, + }, + "enabled": True, + "loggingEnabled": False, + "index": 0, + "metadata": {}, + "description": None, + } + rule: RuleConfig = { + "name": "Test Rule", + "source_zone": "WAN", + "dest_zone": "LAN", + "ip_version": "IPV4", + "action": "ALLOW", + "protocol": "UDP", + } + zone_map = {"WAN": "wan-id", "LAN": "lan-id"} + + result = policy_matches_config(existing_policy, rule, zone_map, "1.2.3.4") # type: ignore[arg-type] + assert result is False + + def test_returns_false_different_action(self) -> None: + existing_policy = { + "id": "pol-1", + "name": "Test Rule", + "source": { + "zoneId": "wan-id", + "trafficFilter": { + "type": "IP_ADDRESS", + "ipAddressFilter": { + "type": "IP_ADDRESSES", + "items": [{"type": "IP_ADDRESS", "value": "1.2.3.4"}], + }, + }, + }, + "destination": {"zoneId": "lan-id", "trafficFilter": None}, + "action": {"type": "ALLOW", "allowReturnTraffic": True}, + "ipProtocolScope": {"ipVersion": "IPV4", "protocolFilter": None}, + "enabled": True, + "loggingEnabled": False, + "index": 0, + "metadata": {}, + "description": None, + } + rule: RuleConfig = { + "name": "Test Rule", + "source_zone": "WAN", + "dest_zone": "LAN", + "ip_version": "IPV4", + "action": "BLOCK", + } + zone_map = {"WAN": "wan-id", "LAN": "lan-id"} + + result = policy_matches_config(existing_policy, rule, zone_map, "1.2.3.4") # type: ignore[arg-type] + assert result is False + + def test_returns_false_different_enabled(self) -> None: + existing_policy = { + "id": "pol-1", + "name": "Test Rule", + "source": { + "zoneId": "wan-id", + "trafficFilter": { + "type": "IP_ADDRESS", + "ipAddressFilter": { + "type": "IP_ADDRESSES", + "items": [{"type": "IP_ADDRESS", "value": "1.2.3.4"}], + }, + }, + }, + "destination": {"zoneId": "lan-id", "trafficFilter": None}, + "action": {"type": "ALLOW", "allowReturnTraffic": True}, + "ipProtocolScope": {"ipVersion": "IPV4", "protocolFilter": None}, + "enabled": False, + "loggingEnabled": False, + "index": 0, + "metadata": {}, + "description": None, + } + rule: RuleConfig = { + "name": "Test Rule", + "source_zone": "WAN", + "dest_zone": "LAN", + "ip_version": "IPV4", + "action": "ALLOW", + "enabled": True, + } + zone_map = {"WAN": "wan-id", "LAN": "lan-id"} + + result = policy_matches_config(existing_policy, rule, zone_map, "1.2.3.4") # type: ignore[arg-type] + assert result is False diff --git a/tests/test_unifi_firewall.py b/tests/test_unifi_firewall.py index 25a6228..32ba4c8 100644 --- a/tests/test_unifi_firewall.py +++ b/tests/test_unifi_firewall.py @@ -9,7 +9,10 @@ from unifi_firewall import ( FirewallPolicyCreatePayload, build_policy_payload, create_policy, + get_action_from_policy, + get_dest_ports_from_policy, get_policy, + get_protocol_from_policy, get_session, get_source_ip_from_policy, list_policies, @@ -297,3 +300,126 @@ class TestGetSourceIpFromPolicy: ip = get_source_ip_from_policy(policy) # type: ignore[arg-type] assert ip is None + + +class TestGetDestPortsFromPolicy: + def test_extracts_ports_and_ranges(self) -> None: + policy: dict[str, Any] = { + "id": "pol-1", + "name": "Test", + "source": {"zoneId": "wan-zone", "trafficFilter": None}, + "destination": { + "zoneId": "lan-zone", + "trafficFilter": { + "type": "PORT", + "portFilter": { + "type": "PORTS", + "matchOpposite": False, + "items": [ + {"type": "PORT_NUMBER", "value": 80}, + {"type": "PORT_NUMBER", "value": 443}, + {"type": "PORT_NUMBER_RANGE", "start": 8000, "stop": 9000}, + ], + }, + }, + }, + "action": {"type": "ALLOW", "allowReturnTraffic": True}, + "ipProtocolScope": {"ipVersion": "IPV4"}, + "enabled": True, + "loggingEnabled": False, + "index": 0, + "metadata": {}, + } + + result = get_dest_ports_from_policy(policy) # type: ignore[arg-type] + assert result is not None + dest_ports, dest_port_ranges = result + assert sorted(dest_ports) == [80, 443] + assert len(dest_port_ranges) == 1 + assert dest_port_ranges[0] == {"start": 8000, "stop": 9000} + + def test_returns_none_no_port_filter(self) -> None: + policy: dict[str, Any] = { + "destination": { + "zoneId": "lan-zone", + "trafficFilter": None, + }, + } + + result = get_dest_ports_from_policy(policy) # type: ignore[arg-type] + assert result is None + + def test_returns_none_empty_items(self) -> None: + policy: dict[str, Any] = { + "destination": { + "zoneId": "lan-zone", + "trafficFilter": { + "type": "PORT", + "portFilter": {"type": "PORTS", "items": []}, + }, + }, + } + + result = get_dest_ports_from_policy(policy) # type: ignore[arg-type] + assert result is None + + +class TestGetProtocolFromPolicy: + def test_extracts_protocol(self) -> None: + policy: dict[str, Any] = { + "ipProtocolScope": { + "ipVersion": "IPV4", + "protocolFilter": { + "type": "NAMED_PROTOCOL", + "protocol": {"name": "TCP"}, + }, + }, + } + + protocol = get_protocol_from_policy(policy) # type: ignore[arg-type] + assert protocol == "TCP" + + def test_returns_none_no_protocol_filter(self) -> None: + policy: dict[str, Any] = { + "ipProtocolScope": { + "ipVersion": "IPV4", + "protocolFilter": None, + }, + } + + protocol = get_protocol_from_policy(policy) # type: ignore[arg-type] + assert protocol is None + + +class TestGetActionFromPolicy: + def test_extracts_allow_action(self) -> None: + policy: dict[str, Any] = { + "action": {"type": "ALLOW", "allowReturnTraffic": True}, + } + + result = get_action_from_policy(policy) # type: ignore[arg-type] + assert result == ("ALLOW", True) + + def test_extracts_allow_action_no_return_traffic(self) -> None: + policy: dict[str, Any] = { + "action": {"type": "ALLOW", "allowReturnTraffic": False}, + } + + result = get_action_from_policy(policy) # type: ignore[arg-type] + assert result == ("ALLOW", False) + + def test_extracts_block_action(self) -> None: + policy: dict[str, Any] = { + "action": {"type": "BLOCK"}, + } + + result = get_action_from_policy(policy) # type: ignore[arg-type] + assert result == ("BLOCK", False) + + def test_extracts_reject_action(self) -> None: + policy: dict[str, Any] = { + "action": {"type": "REJECT"}, + } + + result = get_action_from_policy(policy) # type: ignore[arg-type] + assert result == ("REJECT", False) diff --git a/unifi_firewall.py b/unifi_firewall.py index 6621f4a..b7f19b5 100644 --- a/unifi_firewall.py +++ b/unifi_firewall.py @@ -332,3 +332,87 @@ def get_source_ip_from_policy(policy: FirewallPolicy) -> str | None: value: str | None = first_item.get("value") if isinstance(first_item.get("value"), str) else None return value + + +def get_dest_ports_from_policy(policy: FirewallPolicy) -> tuple[list[int], list[dict[str, Any]]] | None: + """Extract destination port numbers and ranges from a firewall policy. + + Returns (dest_ports_list, dest_port_ranges_list) or None if no port filter exists. + """ + destination = policy.get("destination") + if not destination: + return None + + traffic_filter = destination.get("trafficFilter") + if not traffic_filter: + return None + + port_filter: dict[str, Any] | None = traffic_filter.get("portFilter") + if not port_filter: + return None + + items: list[Any] | None = port_filter.get("items") + if not items: + return None + + dest_ports: list[int] = [] + dest_port_ranges: list[dict[str, Any]] = [] + + for item in items: + item_dict = cast(dict[str, Any], item) if isinstance(item, dict) else None + if not item_dict: + continue + + item_type = item_dict.get("type") + if not isinstance(item_type, str): + continue + + if item_type == "PORT_NUMBER": + value = item_dict.get("value") + if isinstance(value, int): + dest_ports.append(value) + elif item_type == "PORT_NUMBER_RANGE": + start = item_dict.get("start") + stop = item_dict.get("stop") + if isinstance(start, int) and isinstance(stop, int): + dest_port_ranges.append({"start": start, "stop": stop}) + + return (dest_ports, dest_port_ranges) + + +def get_protocol_from_policy(policy: FirewallPolicy) -> str | None: + """Extract protocol name from a firewall policy.""" + ip_protocol_scope = policy.get("ipProtocolScope") + if not ip_protocol_scope: + return None + + protocol_filter: dict[str, Any] | None = ip_protocol_scope.get("protocolFilter") + if not protocol_filter: + return None + + protocol: dict[str, Any] | None = protocol_filter.get("protocol") + if not protocol: + return None + + name: str | None = protocol.get("name") if isinstance(protocol.get("name"), str) else None + return name + + +def get_action_from_policy(policy: FirewallPolicy) -> tuple[str, bool] | None: + """Extract action type and allowReturnTraffic from a firewall policy. + + Returns (action_type, allow_return_traffic) or None if action is missing. + """ + action = policy.get("action") + if not action: + return None + + action_type = action.get("type") + if not action_type: + return None + + allow_return_traffic: bool = False + if action_type == "ALLOW": + allow_return_traffic = bool(action.get("allowReturnTraffic", False)) + + return (action_type, allow_return_traffic)