init

tests/test_unifi_firewall.py

@@ -0,0 +1,299 @@
+"""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_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