426 lines
15 KiB
Python
426 lines
15 KiB
Python
"""Tests for unifi_firewall module."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
from unittest.mock import MagicMock
|
|
|
|
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,
|
|
list_zones,
|
|
update_policy,
|
|
)
|
|
|
|
|
|
class MockHttpError(Exception):
|
|
"""Mock HTTP error for testing."""
|
|
|
|
|
|
class MockResponse:
|
|
def __init__(self, json_data: Any, status_code: int = 200) -> None:
|
|
self.json_data = json_data
|
|
self.status_code = status_code
|
|
|
|
def json(self) -> Any:
|
|
return self.json_data
|
|
|
|
def raise_for_status(self) -> None:
|
|
if self.status_code >= 400:
|
|
raise MockHttpError(f"HTTP {self.status_code}")
|
|
|
|
|
|
class TestGetSession:
|
|
def test_creates_session_with_headers(self) -> None:
|
|
session = get_session("https://unifi.local", "test-token", True)
|
|
|
|
assert session.headers["X-API-Key"] == "test-token"
|
|
assert session.headers["Content-Type"] == "application/json"
|
|
assert session.verify is True
|
|
|
|
def test_verify_ssl_false(self) -> None:
|
|
session = get_session("https://unifi.local", "test-token", False)
|
|
|
|
assert session.verify is False
|
|
|
|
|
|
class TestListZones:
|
|
def test_returns_zones_from_data_key(self) -> None:
|
|
mock_session = MagicMock()
|
|
mock_session.get.return_value = MockResponse({"data": [{"id": "zone-1", "name": "WAN", "networkIds": [], "metadata": {}}]})
|
|
|
|
zones = list_zones(mock_session, "https://unifi.local", "site-1")
|
|
|
|
assert len(zones) == 1
|
|
assert zones[0]["name"] == "WAN"
|
|
mock_session.get.assert_called_once_with("https://unifi.local/proxy/network/integration/v1/sites/site-1/firewall/zones")
|
|
|
|
def test_returns_zones_as_list(self) -> None:
|
|
mock_session = MagicMock()
|
|
mock_session.get.return_value = MockResponse([{"id": "zone-1", "name": "LAN", "networkIds": [], "metadata": {}}])
|
|
|
|
zones = list_zones(mock_session, "https://unifi.local", "site-1")
|
|
|
|
assert len(zones) == 1
|
|
assert zones[0]["name"] == "LAN"
|
|
|
|
|
|
class TestListPolicies:
|
|
def test_returns_policies_without_filter(self) -> None:
|
|
mock_session = MagicMock()
|
|
mock_session.get.return_value = MockResponse({"data": []})
|
|
|
|
policies = list_policies(mock_session, "https://unifi.local", "site-1")
|
|
|
|
assert policies == []
|
|
mock_session.get.assert_called_once_with("https://unifi.local/proxy/network/integration/v1/sites/site-1/firewall/policies", params={})
|
|
|
|
def test_returns_policies_with_name_filter(self) -> None:
|
|
mock_session = MagicMock()
|
|
mock_session.get.return_value = MockResponse({"data": []})
|
|
|
|
policies = list_policies(mock_session, "https://unifi.local", "site-1", name_filter="Test Rule")
|
|
|
|
assert policies == []
|
|
call_args = mock_session.get.call_args
|
|
assert call_args[1]["params"]["filter"] == "name.like('Test Rule')"
|
|
|
|
|
|
class TestGetPolicy:
|
|
def test_returns_single_policy(self) -> None:
|
|
mock_session = MagicMock()
|
|
policy_data = {"id": "pol-1", "name": "Test", "source": {"zoneId": "z1", "trafficFilter": None}, "destination": {"zoneId": "z2", "trafficFilter": None}, "action": {"type": "ALLOW", "allowReturnTraffic": True}, "ipProtocolScope": {"ipVersion": "IPV4"}, "enabled": True, "loggingEnabled": False, "index": 0, "metadata": {}}
|
|
mock_session.get.return_value = MockResponse(policy_data)
|
|
|
|
policy = get_policy(mock_session, "https://unifi.local", "site-1", "pol-1")
|
|
|
|
assert policy["id"] == "pol-1"
|
|
mock_session.get.assert_called_once_with("https://unifi.local/proxy/network/integration/v1/sites/site-1/firewall/policies/pol-1")
|
|
|
|
|
|
class TestCreatePolicy:
|
|
def test_creates_policy(self) -> None:
|
|
mock_session = MagicMock()
|
|
result = {"id": "new-pol", "name": "New Rule", "source": {"zoneId": "z1", "trafficFilter": None}, "destination": {"zoneId": "z2", "trafficFilter": None}, "action": {"type": "ALLOW", "allowReturnTraffic": True}, "ipProtocolScope": {"ipVersion": "IPV4"}, "enabled": True, "loggingEnabled": False, "index": 0, "metadata": {}}
|
|
mock_session.post.return_value = MockResponse(result, status_code=201)
|
|
|
|
payload: FirewallPolicyCreatePayload = build_policy_payload(
|
|
name="New Rule",
|
|
source_ip="1.2.3.4",
|
|
ip_version="IPV4",
|
|
src_zone_id="z1",
|
|
dst_zone_id="z2",
|
|
action_type="ALLOW",
|
|
)
|
|
policy = create_policy(mock_session, "https://unifi.local", "site-1", payload)
|
|
|
|
assert policy["id"] == "new-pol"
|
|
mock_session.post.assert_called_once()
|
|
|
|
|
|
class TestUpdatePolicy:
|
|
def test_updates_policy(self) -> None:
|
|
mock_session = MagicMock()
|
|
result = {"id": "pol-1", "name": "Updated Rule", "source": {"zoneId": "z1", "trafficFilter": None}, "destination": {"zoneId": "z2", "trafficFilter": None}, "action": {"type": "ALLOW", "allowReturnTraffic": True}, "ipProtocolScope": {"ipVersion": "IPV4"}, "enabled": True, "loggingEnabled": False, "index": 0, "metadata": {}}
|
|
mock_session.put.return_value = MockResponse(result)
|
|
|
|
payload: FirewallPolicyCreatePayload = build_policy_payload(
|
|
name="Updated Rule",
|
|
source_ip="5.6.7.8",
|
|
ip_version="IPV4",
|
|
src_zone_id="z1",
|
|
dst_zone_id="z2",
|
|
action_type="ALLOW",
|
|
)
|
|
policy = update_policy(mock_session, "https://unifi.local", "site-1", "pol-1", payload)
|
|
|
|
assert policy["id"] == "pol-1"
|
|
mock_session.put.assert_called_once()
|
|
|
|
|
|
class TestBuildPolicyPayload:
|
|
def test_basic_allow_ipv4(self) -> None:
|
|
payload = build_policy_payload(
|
|
name="Test Rule",
|
|
source_ip="1.2.3.4",
|
|
ip_version="IPV4",
|
|
src_zone_id="wan-zone",
|
|
dst_zone_id="lan-zone",
|
|
action_type="ALLOW",
|
|
)
|
|
|
|
assert payload["name"] == "Test Rule"
|
|
assert payload["source"]["zoneId"] == "wan-zone"
|
|
assert payload["destination"]["zoneId"] == "lan-zone"
|
|
assert payload["action"]["type"] == "ALLOW"
|
|
assert payload["action"]["allowReturnTraffic"] is True
|
|
assert payload["ipProtocolScope"]["ipVersion"] == "IPV4"
|
|
assert payload["enabled"] is True
|
|
assert payload["loggingEnabled"] is False
|
|
|
|
# Check source IP filter
|
|
tf = payload["source"]["trafficFilter"]
|
|
assert tf is not None
|
|
assert tf["type"] == "IP_ADDRESS"
|
|
assert tf["ipAddressFilter"]["type"] == "IP_ADDRESSES"
|
|
assert tf["ipAddressFilter"]["items"][0]["value"] == "1.2.3.4"
|
|
|
|
def test_with_dest_ports(self) -> None:
|
|
payload = build_policy_payload(
|
|
name="Test Rule",
|
|
source_ip="1.2.3.4",
|
|
ip_version="IPV4",
|
|
src_zone_id="wan-zone",
|
|
dst_zone_id="lan-zone",
|
|
action_type="ALLOW",
|
|
dest_ports=[22, 80, 443],
|
|
)
|
|
|
|
dest_tf = payload["destination"]["trafficFilter"]
|
|
assert dest_tf is not None
|
|
assert dest_tf["type"] == "PORT"
|
|
port_filter = dest_tf["portFilter"]
|
|
assert port_filter is not None
|
|
assert port_filter["type"] == "PORTS"
|
|
assert len(port_filter["items"]) == 3
|
|
assert port_filter["items"][0] == {"type": "PORT_NUMBER", "value": 22}
|
|
|
|
def test_with_port_ranges(self) -> None:
|
|
payload = build_policy_payload(
|
|
name="Test Rule",
|
|
source_ip="1.2.3.4",
|
|
ip_version="IPV4",
|
|
src_zone_id="wan-zone",
|
|
dst_zone_id="lan-zone",
|
|
action_type="ALLOW",
|
|
dest_port_ranges=[{"start": 8000, "stop": 9000}],
|
|
)
|
|
|
|
dest_tf = payload["destination"]["trafficFilter"]
|
|
assert dest_tf is not None
|
|
assert dest_tf["type"] == "PORT"
|
|
port_filter = dest_tf["portFilter"]
|
|
assert port_filter is not None
|
|
assert port_filter["items"][0] == {"type": "PORT_NUMBER_RANGE", "start": 8000, "stop": 9000}
|
|
|
|
def test_with_protocol(self) -> None:
|
|
payload = build_policy_payload(
|
|
name="Test Rule",
|
|
source_ip="1.2.3.4",
|
|
ip_version="IPV6",
|
|
src_zone_id="wan-zone",
|
|
dst_zone_id="lan-zone",
|
|
action_type="ALLOW",
|
|
protocol="tcp",
|
|
)
|
|
|
|
assert payload["ipProtocolScope"]["ipVersion"] == "IPV6"
|
|
pf = payload["ipProtocolScope"]["protocolFilter"]
|
|
assert pf is not None
|
|
assert pf["protocol"] == {"name": "TCP"}
|
|
|
|
def test_block_action(self) -> None:
|
|
payload = build_policy_payload(
|
|
name="Block Rule",
|
|
source_ip="5.6.7.8",
|
|
ip_version="IPV4",
|
|
src_zone_id="wan-zone",
|
|
dst_zone_id="lan-zone",
|
|
action_type="BLOCK",
|
|
)
|
|
|
|
assert payload["action"]["type"] == "BLOCK"
|
|
assert "allowReturnTraffic" not in payload["action"]
|
|
|
|
|
|
class TestGetSourceIpFromPolicy:
|
|
def test_extracts_ip_from_policy(self) -> None:
|
|
policy: dict[str, Any] = {
|
|
"id": "pol-1",
|
|
"name": "Test",
|
|
"source": {
|
|
"zoneId": "wan-zone",
|
|
"trafficFilter": {
|
|
"type": "IP_ADDRESS",
|
|
"ipAddressFilter": {
|
|
"type": "IP_ADDRESSES",
|
|
"items": [{"type": "IP_ADDRESS", "value": "1.2.3.4"}],
|
|
},
|
|
},
|
|
},
|
|
"destination": {"zoneId": "lan-zone", "trafficFilter": None},
|
|
"action": {"type": "ALLOW", "allowReturnTraffic": True},
|
|
"ipProtocolScope": {"ipVersion": "IPV4"},
|
|
"enabled": True,
|
|
"loggingEnabled": False,
|
|
"index": 0,
|
|
"metadata": {},
|
|
}
|
|
|
|
ip = get_source_ip_from_policy(policy) # type: ignore[arg-type]
|
|
assert ip == "1.2.3.4"
|
|
|
|
def test_returns_none_no_traffic_filter(self) -> None:
|
|
policy: dict[str, Any] = {
|
|
"source": {"zoneId": "wan-zone", "trafficFilter": None},
|
|
}
|
|
|
|
ip = get_source_ip_from_policy(policy) # type: ignore[arg-type]
|
|
assert ip is None
|
|
|
|
def test_returns_none_wrong_filter_type(self) -> None:
|
|
policy: dict[str, Any] = {
|
|
"source": {
|
|
"zoneId": "wan-zone",
|
|
"trafficFilter": {"type": "NETWORK"},
|
|
},
|
|
}
|
|
|
|
ip = get_source_ip_from_policy(policy) # type: ignore[arg-type]
|
|
assert ip is None
|
|
|
|
def test_returns_none_empty_items(self) -> None:
|
|
policy: dict[str, Any] = {
|
|
"source": {
|
|
"zoneId": "wan-zone",
|
|
"trafficFilter": {
|
|
"type": "IP_ADDRESS",
|
|
"ipAddressFilter": {"type": "IP_ADDRESSES", "items": []},
|
|
},
|
|
},
|
|
}
|
|
|
|
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)
|