diff --git a/homeassistant/components/unifi/hub/entity_loader.py b/homeassistant/components/unifi/hub/entity_loader.py
index 64403152b0c234c9169a5badbaeac6abae6ff910..84948a92e98819b61afbf1c6ced2812df5673f36 100644
--- a/homeassistant/components/unifi/hub/entity_loader.py
+++ b/homeassistant/components/unifi/hub/entity_loader.py
@@ -46,6 +46,7 @@ class UnifiEntityLoader:
hub.api.port_forwarding.update,
hub.api.sites.update,
hub.api.system_information.update,
+ hub.api.firewall_policies.update,
hub.api.traffic_rules.update,
hub.api.traffic_routes.update,
hub.api.wlans.update,
diff --git a/homeassistant/components/unifi/icons.json b/homeassistant/components/unifi/icons.json
index 6874bb5ae03d89133176dd3499ad4efbdc6f8063..616d7cb185fdd7134dfbe286b93f550859d9f87c 100644
--- a/homeassistant/components/unifi/icons.json
+++ b/homeassistant/components/unifi/icons.json
@@ -55,6 +55,9 @@
"off": "mdi:network-off"
}
},
+ "firewall_policy_control": {
+ "default": "mdi:security-network"
+ },
"port_forward_control": {
"default": "mdi:upload-network"
},
diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py
index de0e8d3f4126edddb08392900fe5655cbbd05094..282d0c9ae9378b9df63750a7c74fb43bef8a33cb 100644
--- a/homeassistant/components/unifi/switch.py
+++ b/homeassistant/components/unifi/switch.py
@@ -4,6 +4,7 @@ Support for controlling power supply of clients which are powered over Ethernet
Support for controlling network access of clients selected in option flow.
Support for controlling deep packet inspection (DPI) restriction groups.
Support for controlling WLAN availability.
+Support for controlling zone based traffic rules.
"""
from __future__ import annotations
@@ -17,6 +18,7 @@ import aiounifi
from aiounifi.interfaces.api_handlers import ItemEvent
from aiounifi.interfaces.clients import Clients
from aiounifi.interfaces.dpi_restriction_groups import DPIRestrictionGroups
+from aiounifi.interfaces.firewall_policies import FirewallPolicies
from aiounifi.interfaces.outlets import Outlets
from aiounifi.interfaces.port_forwarding import PortForwarding
from aiounifi.interfaces.ports import Ports
@@ -29,6 +31,7 @@ from aiounifi.models.device import DeviceSetOutletRelayRequest
from aiounifi.models.dpi_restriction_app import DPIRestrictionAppEnableRequest
from aiounifi.models.dpi_restriction_group import DPIRestrictionGroup
from aiounifi.models.event import Event, EventKey
+from aiounifi.models.firewall_policy import FirewallPolicy, FirewallPolicyUpdateRequest
from aiounifi.models.outlet import Outlet
from aiounifi.models.port import Port
from aiounifi.models.port_forward import PortForward, PortForwardEnableRequest
@@ -129,6 +132,24 @@ async def async_dpi_group_control_fn(hub: UnifiHub, obj_id: str, target: bool) -
)
+async def async_firewall_policy_control_fn(
+ hub: UnifiHub, obj_id: str, target: bool
+) -> None:
+ """Control firewall policy state."""
+ policy = hub.api.firewall_policies[obj_id].raw
+ policy["enabled"] = target
+ await hub.api.request(FirewallPolicyUpdateRequest.create(policy))
+ # Update the policies so the UI is updated appropriately
+ await hub.api.firewall_policies.update()
+
+
+@callback
+def async_firewall_policy_supported_fn(hub: UnifiHub, obj_id: str) -> bool:
+ """Check if firewall policy is able to be controlled. Predefined policies are unable to be turned off."""
+ policy = hub.api.firewall_policies[obj_id]
+ return not policy.predefined
+
+
@callback
def async_outlet_switching_supported_fn(hub: UnifiHub, obj_id: str) -> bool:
"""Determine if an outlet supports switching."""
@@ -236,6 +257,20 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
supported_fn=lambda hub, obj_id: bool(hub.api.dpi_groups[obj_id].dpiapp_ids),
unique_id_fn=lambda hub, obj_id: obj_id,
),
+ UnifiSwitchEntityDescription[FirewallPolicies, FirewallPolicy](
+ key="Firewall policy control",
+ translation_key="firewall_policy_control",
+ device_class=SwitchDeviceClass.SWITCH,
+ entity_category=EntityCategory.CONFIG,
+ api_handler_fn=lambda api: api.firewall_policies,
+ control_fn=async_firewall_policy_control_fn,
+ device_info_fn=async_unifi_network_device_info_fn,
+ is_on_fn=lambda hub, firewall_policy: firewall_policy.enabled,
+ name_fn=lambda firewall_policy: firewall_policy.name,
+ object_fn=lambda api, obj_id: api.firewall_policies[obj_id],
+ unique_id_fn=lambda hub, obj_id: f"firewall_policy-{obj_id}",
+ supported_fn=async_firewall_policy_supported_fn,
+ ),
UnifiSwitchEntityDescription[Outlets, Outlet](
key="Outlet control",
device_class=SwitchDeviceClass.OUTLET,
diff --git a/tests/components/unifi/conftest.py b/tests/components/unifi/conftest.py
index ec7a05957319efd161ce9bf9f106358dd9bddce5..4075aa0ad59945bd24c5f4925f99c203a18c3cd4 100644
--- a/tests/components/unifi/conftest.py
+++ b/tests/components/unifi/conftest.py
@@ -172,6 +172,7 @@ def fixture_request(
device_payload: list[dict[str, Any]],
dpi_app_payload: list[dict[str, Any]],
dpi_group_payload: list[dict[str, Any]],
+ firewall_policy_payload: list[dict[str, Any]],
port_forward_payload: list[dict[str, Any]],
traffic_rule_payload: list[dict[str, Any]],
traffic_route_payload: list[dict[str, Any]],
@@ -211,6 +212,9 @@ def fixture_request(
mock_get_request(f"/api/s/{site_id}/stat/device", device_payload)
mock_get_request(f"/api/s/{site_id}/rest/dpiapp", dpi_app_payload)
mock_get_request(f"/api/s/{site_id}/rest/dpigroup", dpi_group_payload)
+ mock_get_request(
+ f"/v2/api/site/{site_id}/firewall-policies", firewall_policy_payload
+ )
mock_get_request(f"/api/s/{site_id}/rest/portforward", port_forward_payload)
mock_get_request(f"/api/s/{site_id}/stat/sysinfo", system_information_payload)
mock_get_request(f"/api/s/{site_id}/rest/wlanconf", wlan_payload)
@@ -253,6 +257,12 @@ def fixture_dpi_group_data() -> list[dict[str, Any]]:
return []
+@pytest.fixture(name="firewall_policy_payload")
+def firewall_policy_payload_data() -> list[dict[str, Any]]:
+ """Firewall policy data."""
+ return []
+
+
@pytest.fixture(name="port_forward_payload")
def fixture_port_forward_data() -> list[dict[str, Any]]:
"""Port forward data."""
diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py
index e4765d1181e6aff4386e77152ba65f0775efff58..c8ee786895cc48fd9b61b9eb12979bd9255ac36c 100644
--- a/tests/components/unifi/test_switch.py
+++ b/tests/components/unifi/test_switch.py
@@ -827,6 +827,45 @@ TRAFFIC_ROUTE = {
],
}
+FIREWALL_POLICY = {
+ "_id": "678ceb9fe3849d293243405c",
+ "action": "ALLOW",
+ "connection_state_type": "ALL",
+ "connection_states": [],
+ "create_allow_respond": True,
+ "description": "",
+ "destination": {
+ "match_opposite_ports": False,
+ "matching_target": "ANY",
+ "port_matching_type": "ANY",
+ "zone_id": "678ccc26e3849d2932432e26",
+ },
+ "enabled": True,
+ "icmp_typename": "ANY",
+ "icmp_v6_typename": "ANY",
+ "index": 10000,
+ "ip_version": "BOTH",
+ "logging": False,
+ "match_ip_sec": False,
+ "match_opposite_protocol": False,
+ "name": "Allow internal to IoT",
+ "predefined": False,
+ "protocol": "all",
+ "schedule": {
+ "mode": "EVERY_DAY",
+ "repeat_on_days": [],
+ "time_all_day": False,
+ "time_range_end": "12:00",
+ "time_range_start": "09:00",
+ },
+ "source": {
+ "match_opposite_ports": False,
+ "matching_target": "ANY",
+ "port_matching_type": "ANY",
+ "zone_id": "678c63bc2d97692f08adcdfa",
+ },
+}
+
@pytest.mark.parametrize(
"config_entry_options", [{CONF_BLOCK_CLIENT: [BLOCKED["mac"]]}]
@@ -1226,6 +1265,62 @@ async def test_traffic_routes(
assert aioclient_mock.mock_calls[call_count][2] == expected_enable_call
+@pytest.mark.parametrize(("firewall_policy_payload"), [([FIREWALL_POLICY])])
+async def test_firewall_policies(
+ hass: HomeAssistant,
+ aioclient_mock: AiohttpClientMocker,
+ config_entry_setup: MockConfigEntry,
+ firewall_policy_payload: list[dict[str, Any]],
+) -> None:
+ """Test control of UniFi firewall policies."""
+ assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
+
+ # Validate state object
+ assert (
+ hass.states.get("switch.unifi_network_allow_internal_to_iot").state == STATE_ON
+ )
+
+ firewall_policy = deepcopy(firewall_policy_payload[0])
+
+ # Disable firewall policy
+ aioclient_mock.put(
+ f"https://{config_entry_setup.data[CONF_HOST]}:1234"
+ f"/v2/api/site/{config_entry_setup.data[CONF_SITE_ID]}"
+ f"/firewall-policies/{firewall_policy['_id']}",
+ )
+
+ call_count = aioclient_mock.call_count
+
+ await hass.services.async_call(
+ SWITCH_DOMAIN,
+ "turn_off",
+ {"entity_id": "switch.unifi_network_allow_internal_to_iot"},
+ blocking=True,
+ )
+ # Updating the value for firewall policies will make another call to retrieve the values
+ assert aioclient_mock.call_count == call_count + 2
+ expected_disable_call = deepcopy(firewall_policy)
+ expected_disable_call["enabled"] = False
+
+ assert aioclient_mock.mock_calls[call_count][2] == expected_disable_call
+
+ call_count = aioclient_mock.call_count
+
+ # Enable firewall policy
+ await hass.services.async_call(
+ SWITCH_DOMAIN,
+ "turn_on",
+ {"entity_id": "switch.unifi_network_allow_internal_to_iot"},
+ blocking=True,
+ )
+
+ expected_enable_call = deepcopy(firewall_policy)
+ expected_enable_call["enabled"] = True
+
+ assert aioclient_mock.call_count == call_count + 2
+ assert aioclient_mock.mock_calls[call_count][2] == expected_enable_call
+
+
@pytest.mark.parametrize(
("device_payload", "entity_id", "outlet_index", "expected_switches"),
[
@@ -1677,6 +1772,7 @@ async def test_updating_unique_id(
@pytest.mark.parametrize("dpi_group_payload", [DPI_GROUPS])
@pytest.mark.parametrize("port_forward_payload", [[PORT_FORWARD_PLEX]])
@pytest.mark.parametrize(("traffic_rule_payload"), [([TRAFFIC_RULE])])
+@pytest.mark.parametrize("firewall_policy_payload", [[FIREWALL_POLICY]])
@pytest.mark.parametrize("wlan_payload", [[WLAN]])
@pytest.mark.usefixtures("config_entry_setup")
@pytest.mark.usefixtures("entity_registry_enabled_by_default")
@@ -1691,6 +1787,7 @@ async def test_hub_state_change(
"switch.block_media_streaming",
"switch.unifi_network_plex",
"switch.unifi_network_test_traffic_rule",
+ "switch.unifi_network_allow_internal_to_iot",
"switch.ssid_1",
)
for entity_id in entity_ids: