512 lines
18 KiB
Python
512 lines
18 KiB
Python
"""Tests for main module."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from main import (
|
|
RuleConfig,
|
|
build_zone_map,
|
|
load_config,
|
|
policy_matches_config,
|
|
process_rule,
|
|
send_ntfy_notification,
|
|
)
|
|
from unifi_firewall import FirewallZone
|
|
|
|
|
|
class TestBuildZoneMap:
|
|
def test_builds_map_from_zones(self) -> None:
|
|
zones: list[FirewallZone] = [
|
|
{"id": "zone-wan", "name": "WAN", "networkIds": [], "metadata": {}},
|
|
{"id": "zone-lan", "name": "LAN", "networkIds": [], "metadata": {}},
|
|
]
|
|
|
|
zone_map = build_zone_map(zones)
|
|
|
|
assert zone_map == {"WAN": "zone-wan", "LAN": "zone-lan"}
|
|
|
|
def test_empty_zones(self) -> None:
|
|
zone_map = build_zone_map([])
|
|
assert zone_map == {}
|
|
|
|
|
|
class TestLoadConfig:
|
|
def test_loads_valid_config(self, tmp_path: Path) -> None:
|
|
config_file = tmp_path / "rules.yaml"
|
|
config_file.write_text(
|
|
"""
|
|
rules:
|
|
- name: "Test Rule"
|
|
source_zone: "WAN"
|
|
dest_zone: "LAN"
|
|
ip_version: "IPV4"
|
|
action: "ALLOW"
|
|
"""
|
|
)
|
|
|
|
rules = load_config(str(config_file))
|
|
|
|
assert len(rules) == 1
|
|
assert rules[0].get("name") == "Test Rule"
|
|
|
|
def test_exits_on_missing_file(self) -> None:
|
|
with patch("sys.exit") as mock_exit:
|
|
load_config("/nonexistent/rules.yaml")
|
|
mock_exit.assert_called_once_with(1)
|
|
|
|
def test_exits_on_invalid_yaml(self, tmp_path: Path) -> None:
|
|
config_file = tmp_path / "rules.yaml"
|
|
config_file.write_text("invalid: yaml: content: [")
|
|
|
|
with patch("sys.exit") as mock_exit:
|
|
load_config(str(config_file))
|
|
mock_exit.assert_called_once_with(1)
|
|
|
|
|
|
class TestSendNtfyNotification:
|
|
def test_sends_notification(self) -> None:
|
|
mock_response = MagicMock()
|
|
mock_response.raise_for_status = MagicMock()
|
|
|
|
with (
|
|
patch("main.NTFY_URL", "https://ntfy.sh"),
|
|
patch("main.NTFY_TOPIC", "test"),
|
|
patch("main.NTFY_API_KEY", "test-key"),
|
|
patch("main.requests.post", return_value=mock_response) as mock_post,
|
|
):
|
|
send_ntfy_notification("Title", "Message", priority=4)
|
|
|
|
mock_post.assert_called_once()
|
|
call_kwargs = mock_post.call_args
|
|
assert call_kwargs[0][0] == "https://ntfy.sh/test"
|
|
assert call_kwargs[1]["data"] == b"Message"
|
|
assert call_kwargs[1]["headers"]["Title"] == "Title"
|
|
assert call_kwargs[1]["headers"]["Priority"] == "4"
|
|
|
|
def test_skips_when_no_url(self) -> None:
|
|
with (
|
|
patch("main.NTFY_URL", ""),
|
|
patch("main.requests.post") as mock_post,
|
|
):
|
|
send_ntfy_notification("Title", "Message")
|
|
mock_post.assert_not_called()
|
|
|
|
def test_handles_request_error(self) -> None:
|
|
import requests as req
|
|
|
|
with (
|
|
patch("main.NTFY_URL", "https://ntfy.sh"),
|
|
patch("main.NTFY_TOPIC", "test"),
|
|
patch("main.requests.post", side_effect=req.RequestException("timeout")),
|
|
):
|
|
# Should not raise, only log warning
|
|
send_ntfy_notification("Title", "Message")
|
|
|
|
|
|
class TestProcessRule:
|
|
def test_creates_new_rule(self) -> None:
|
|
mock_session = MagicMock()
|
|
mock_session.get.return_value = MagicMock(json=MagicMock(return_value={"data": []}), raise_for_status=MagicMock())
|
|
mock_session.post.return_value = MagicMock(
|
|
json=MagicMock(return_value={"id": "new-id", "name": "Test Rule"}),
|
|
raise_for_status=MagicMock(),
|
|
)
|
|
|
|
rule: RuleConfig = {
|
|
"name": "Test Rule",
|
|
"source_zone": "WAN",
|
|
"dest_zone": "LAN",
|
|
"ip_version": "IPV4",
|
|
"action": "ALLOW",
|
|
}
|
|
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"] == "created"
|
|
assert change["rule_name"] == "AUTO-Test Rule"
|
|
assert change["error"] is None
|
|
|
|
def test_skips_rule_with_matching_config(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": 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())
|
|
|
|
rule: RuleConfig = {
|
|
"name": "Test Rule",
|
|
"source_zone": "WAN",
|
|
"dest_zone": "LAN",
|
|
"ip_version": "IPV4",
|
|
}
|
|
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"] == "skipped"
|
|
mock_session.put.assert_not_called()
|
|
mock_session.post.assert_not_called()
|
|
|
|
def test_updates_rule_with_different_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": "5.6.7.8"}],
|
|
},
|
|
},
|
|
},
|
|
}
|
|
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",
|
|
}
|
|
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()
|
|
|
|
def test_fails_on_missing_source_zone(self) -> None:
|
|
mock_session = MagicMock()
|
|
|
|
rule: RuleConfig = {
|
|
"name": "Test Rule",
|
|
"source_zone": "UNKNOWN",
|
|
"dest_zone": "LAN",
|
|
"ip_version": "IPV4",
|
|
}
|
|
zone_map = {"LAN": "lan-id"}
|
|
|
|
change = process_rule(mock_session, "https://unifi.local", "site-1", rule, zone_map, "1.2.3.4")
|
|
|
|
assert change["action"] == "failed"
|
|
assert "Source zone" in (change["error"] or "")
|
|
|
|
def test_fails_on_api_error(self) -> None:
|
|
mock_session = MagicMock()
|
|
mock_session.get.side_effect = Exception("Connection error")
|
|
|
|
rule: RuleConfig = {
|
|
"name": "Test Rule",
|
|
"source_zone": "WAN",
|
|
"dest_zone": "LAN",
|
|
"ip_version": "IPV4",
|
|
}
|
|
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"] == "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
|