335 lines
9.7 KiB
Python
335 lines
9.7 KiB
Python
"""UniFi Network API client for firewall policy management."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any, Literal, TypedDict, cast
|
|
|
|
import requests
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ========== Response TypedDicts ==========
|
|
|
|
|
|
class FirewallZone(TypedDict):
|
|
"""Firewall zone from GET /firewall/zones."""
|
|
|
|
id: str
|
|
name: str
|
|
networkIds: list[str]
|
|
metadata: dict[str, Any]
|
|
|
|
|
|
class FirewallPolicySource(TypedDict):
|
|
"""Firewall policy source."""
|
|
|
|
zoneId: str
|
|
trafficFilter: dict[str, Any] | None
|
|
|
|
|
|
class FirewallPolicyDestination(TypedDict):
|
|
"""Firewall policy destination."""
|
|
|
|
zoneId: str
|
|
trafficFilter: dict[str, Any] | None
|
|
|
|
|
|
class FirewallPolicyActionAllowDto(TypedDict):
|
|
"""Allow action for firewall policy."""
|
|
|
|
type: Literal["ALLOW"]
|
|
allowReturnTraffic: bool
|
|
|
|
|
|
class FirewallPolicyActionBlockDto(TypedDict):
|
|
"""Block action for firewall policy."""
|
|
|
|
type: Literal["BLOCK"]
|
|
|
|
|
|
class FirewallPolicyActionRejectDto(TypedDict):
|
|
"""Reject action for firewall policy."""
|
|
|
|
type: Literal["REJECT"]
|
|
|
|
|
|
FirewallPolicyAction = FirewallPolicyActionAllowDto | FirewallPolicyActionBlockDto | FirewallPolicyActionRejectDto
|
|
|
|
|
|
class FirewallPolicyIpv4ProtocolScopeDto(TypedDict):
|
|
"""IPv4 protocol scope for firewall policy."""
|
|
|
|
ipVersion: Literal["IPV4"]
|
|
protocolFilter: dict[str, Any] | None
|
|
|
|
|
|
class FirewallPolicyIpv6ProtocolScopeDto(TypedDict):
|
|
"""IPv6 protocol scope for firewall policy."""
|
|
|
|
ipVersion: Literal["IPV6"]
|
|
protocolFilter: dict[str, Any] | None
|
|
|
|
|
|
FirewallPolicyIpProtocolScope = FirewallPolicyIpv4ProtocolScopeDto | FirewallPolicyIpv6ProtocolScopeDto
|
|
|
|
|
|
class FirewallPolicy(TypedDict):
|
|
"""Firewall policy from GET /firewall/policies."""
|
|
|
|
id: str
|
|
name: str
|
|
source: FirewallPolicySource
|
|
destination: FirewallPolicyDestination
|
|
action: FirewallPolicyAction
|
|
ipProtocolScope: FirewallPolicyIpProtocolScope
|
|
enabled: bool
|
|
loggingEnabled: bool
|
|
index: int
|
|
metadata: dict[str, Any]
|
|
description: str | None
|
|
connectionStateFilter: list[str] | None
|
|
ipsecFilter: str | None
|
|
schedule: dict[str, Any] | None
|
|
|
|
|
|
# ========== Request TypedDicts ==========
|
|
|
|
|
|
class FirewallPolicyCreatePayload(TypedDict):
|
|
"""Payload for POST/PUT /firewall/policies."""
|
|
|
|
name: str
|
|
source: FirewallPolicySource
|
|
destination: FirewallPolicyDestination
|
|
action: FirewallPolicyAction
|
|
ipProtocolScope: FirewallPolicyIpProtocolScope
|
|
enabled: bool
|
|
loggingEnabled: bool
|
|
description: str | None
|
|
connectionStateFilter: list[str] | None
|
|
ipsecFilter: str | None
|
|
schedule: dict[str, Any] | None
|
|
|
|
|
|
# ========== Session helper ==========
|
|
|
|
|
|
def get_session(host: str, api_token: str, verify_ssl: bool) -> requests.Session:
|
|
"""Create an authenticated requests session for the UniFi API."""
|
|
session = requests.Session()
|
|
session.headers.update(
|
|
{
|
|
"X-API-Key": api_token,
|
|
"Content-Type": "application/json",
|
|
}
|
|
)
|
|
session.verify = verify_ssl
|
|
return session
|
|
|
|
|
|
# ========== API functions ==========
|
|
|
|
|
|
def _api_url(host: str, site_id: str, path: str) -> str:
|
|
"""Build full API URL."""
|
|
return f"{host}/proxy/network/integration/v1/sites/{site_id}{path}"
|
|
|
|
|
|
def list_zones(session: requests.Session, host: str, site_id: str) -> list[FirewallZone]:
|
|
"""List all firewall zones."""
|
|
url = _api_url(host, site_id, "/firewall/zones")
|
|
logger.debug("Listing firewall zones from %s", url)
|
|
response = session.get(url)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
return data.get("data", data) if isinstance(data, dict) else data # type: ignore[no-any-return]
|
|
|
|
|
|
def list_policies(session: requests.Session, host: str, site_id: str, name_filter: str | None = None) -> list[FirewallPolicy]:
|
|
"""List firewall policies, optionally filtered by name."""
|
|
url = _api_url(host, site_id, "/firewall/policies")
|
|
params: dict[str, str] = {}
|
|
if name_filter:
|
|
params["filter"] = f'name.like(\'{name_filter}\')'
|
|
logger.debug("Listing firewall policies from %s (filter=%s)", url, params.get("filter"))
|
|
response = session.get(url, params=params)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
return data.get("data", data) if isinstance(data, dict) else data # type: ignore[no-any-return]
|
|
|
|
|
|
def get_policy(session: requests.Session, host: str, site_id: str, policy_id: str) -> FirewallPolicy:
|
|
"""Get a single firewall policy by ID."""
|
|
url = _api_url(host, site_id, f"/firewall/policies/{policy_id}")
|
|
logger.debug("Getting firewall policy %s from %s", policy_id, url)
|
|
response = session.get(url)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
|
|
def create_policy(session: requests.Session, host: str, site_id: str, payload: FirewallPolicyCreatePayload) -> FirewallPolicy:
|
|
"""Create a new firewall policy."""
|
|
url = _api_url(host, site_id, "/firewall/policies")
|
|
logger.debug("Creating firewall policy '%s' at %s", payload["name"], url)
|
|
response = session.post(url, json=cast(dict[str, Any], payload))
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
|
|
def update_policy(session: requests.Session, host: str, site_id: str, policy_id: str, payload: FirewallPolicyCreatePayload) -> FirewallPolicy:
|
|
"""Update an existing firewall policy (full replace)."""
|
|
url = _api_url(host, site_id, f"/firewall/policies/{policy_id}")
|
|
logger.debug("Updating firewall policy %s at %s", policy_id, url)
|
|
response = session.put(url, json=cast(dict[str, Any], payload))
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
|
|
# ========== Payload builder ==========
|
|
|
|
|
|
def build_policy_payload(
|
|
name: str,
|
|
source_ip: str,
|
|
ip_version: Literal["IPV4", "IPV6"],
|
|
src_zone_id: str,
|
|
dst_zone_id: str,
|
|
action_type: Literal["ALLOW", "BLOCK", "REJECT"],
|
|
allow_return_traffic: bool = True,
|
|
protocol: str | None = None,
|
|
dest_ports: list[int] | None = None,
|
|
dest_port_ranges: list[dict[str, int]] | None = None,
|
|
logging_enabled: bool = False,
|
|
enabled: bool = True,
|
|
description: str | None = None,
|
|
) -> FirewallPolicyCreatePayload:
|
|
"""Build a full firewall policy payload from rule configuration."""
|
|
# Build source with IP filter
|
|
source: FirewallPolicySource = {
|
|
"zoneId": src_zone_id,
|
|
"trafficFilter": {
|
|
"type": "IP_ADDRESS",
|
|
"ipAddressFilter": {
|
|
"type": "IP_ADDRESSES",
|
|
"matchOpposite": False,
|
|
"items": [{"type": "IP_ADDRESS", "value": source_ip}],
|
|
},
|
|
"portFilter": None,
|
|
},
|
|
}
|
|
|
|
# Build destination with optional port filter
|
|
dest_port_items: list[dict[str, Any]] = []
|
|
if dest_ports:
|
|
dest_port_items.extend({"type": "PORT_NUMBER", "value": p} for p in dest_ports)
|
|
if dest_port_ranges:
|
|
dest_port_items.extend(
|
|
{"type": "PORT_NUMBER_RANGE", "start": r["start"], "stop": r["stop"]} for r in dest_port_ranges
|
|
)
|
|
|
|
dest_port_filter: dict[str, Any] | None = None
|
|
if dest_port_items:
|
|
dest_port_filter = {
|
|
"type": "PORTS",
|
|
"matchOpposite": False,
|
|
"items": dest_port_items,
|
|
}
|
|
|
|
destination: FirewallPolicyDestination = {
|
|
"zoneId": dst_zone_id,
|
|
"trafficFilter": (
|
|
{
|
|
"type": "PORT",
|
|
"portFilter": dest_port_filter,
|
|
}
|
|
if dest_port_filter
|
|
else None
|
|
),
|
|
}
|
|
|
|
# Build action
|
|
if action_type == "ALLOW":
|
|
action: FirewallPolicyAction = {"type": "ALLOW", "allowReturnTraffic": allow_return_traffic}
|
|
elif action_type == "BLOCK":
|
|
action = {"type": "BLOCK"}
|
|
else:
|
|
action = {"type": "REJECT"}
|
|
|
|
# Build protocol scope
|
|
protocol_filter: dict[str, Any] | None = None
|
|
if protocol:
|
|
protocol_filter = {
|
|
"type": "NAMED_PROTOCOL",
|
|
"matchOpposite": False,
|
|
"protocol": {
|
|
"name": protocol.upper(),
|
|
},
|
|
}
|
|
|
|
if ip_version == "IPV4":
|
|
ip_protocol_scope: FirewallPolicyIpProtocolScope = {
|
|
"ipVersion": "IPV4",
|
|
"protocolFilter": protocol_filter,
|
|
}
|
|
else:
|
|
ip_protocol_scope = {
|
|
"ipVersion": "IPV6",
|
|
"protocolFilter": protocol_filter,
|
|
}
|
|
|
|
return {
|
|
"name": name,
|
|
"source": source,
|
|
"destination": destination,
|
|
"action": action,
|
|
"ipProtocolScope": ip_protocol_scope,
|
|
"enabled": enabled,
|
|
"loggingEnabled": logging_enabled,
|
|
"description": description,
|
|
"connectionStateFilter": None,
|
|
"ipsecFilter": None,
|
|
"schedule": None,
|
|
}
|
|
|
|
|
|
# ========== IP extraction ==========
|
|
|
|
|
|
def get_source_ip_from_policy(policy: FirewallPolicy) -> str | None:
|
|
"""Extract the source IP address from a firewall policy."""
|
|
source = policy.get("source")
|
|
if not source:
|
|
return None
|
|
|
|
traffic_filter = source.get("trafficFilter")
|
|
if not traffic_filter:
|
|
return None
|
|
|
|
tf_type: str | None = traffic_filter.get("type")
|
|
if tf_type != "IP_ADDRESS":
|
|
return None
|
|
|
|
ip_filter: dict[str, Any] | None = traffic_filter.get("ipAddressFilter")
|
|
if not ip_filter:
|
|
return None
|
|
|
|
ip_filter_type: str | None = ip_filter.get("type")
|
|
if ip_filter_type != "IP_ADDRESSES":
|
|
return None
|
|
|
|
items: list[Any] | None = ip_filter.get("items")
|
|
if not items or len(items) == 0:
|
|
return None
|
|
|
|
first_item: dict[str, Any] | None = cast(dict[str, Any], items[0]) if isinstance(items[0], dict) else None
|
|
if not first_item:
|
|
return None
|
|
|
|
if first_item.get("type") != "IP_ADDRESS":
|
|
return None
|
|
|
|
value: str | None = first_item.get("value") if isinstance(first_item.get("value"), str) else None
|
|
return value
|