This commit is contained in:
@@ -0,0 +1,334 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user