Feat/port profile switch entities

custom_components/unifi_network_rules/coordinator.py

@@ -8,10 +8,9 @@
 
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
-from homeassistant.helpers.dispatcher import async_dispatcher_connect, async_dispatcher_send
+from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.entity_registry import async_get as async_get_entity_registry
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 
 from aiounifi.models.traffic_route import TrafficRoute
@@ -22,14 +21,16 @@
 from aiounifi.models.wlan import Wlan
 from aiounifi.models.device import Device
 
-from .const import DOMAIN, LOGGER, CONF_UPDATE_INTERVAL, DEFAULT_UPDATE_INTERVAL, DEBUG_WEBSOCKET, LOG_TRIGGERS
+from .const import DOMAIN, LOGGER, DEFAULT_UPDATE_INTERVAL, LOG_TRIGGERS
 from .udm import UDMAPI
-from .websocket import SIGNAL_WEBSOCKET_MESSAGE, UnifiRuleWebsocket
-from .helpers.rule import get_rule_id, get_rule_name, get_rule_enabled, get_child_unique_id
+from .websocket import UnifiRuleWebsocket
+from .helpers.rule import get_rule_id, get_child_unique_id
 from .utils.logger import log_data, log_websocket
 from .models.firewall_rule import FirewallRule
 from .models.qos_rule import QoSRule
 from .models.vpn_config import VPNConfig
+from .models.port_profile import PortProfile
+from .models.network import NetworkConf
 
 # This is a fallback if no update_interval is specified
 SCAN_INTERVAL = timedelta(seconds=60)
@@ -110,6 +111,8 @@ def __init__(
         self.vpn_clients: List[VPNConfig] = []
         self.vpn_servers: List[VPNConfig] = []
         self.devices: List[Device] = []  # For LED toggle switches
+        self.port_profiles: List[PortProfile] = []
+        self.networks: List[NetworkConf] = []
 
         # For dynamic entity creation
         self.async_add_entities_callback: AddEntitiesCallback | None = None
@@ -370,6 +373,8 @@ async def _async_update_data(self) -> Dict[str, List[Any]]:
                     "vpn_clients": [],
                     "vpn_servers": [],
                     "devices": [],
+                    "port_profiles": [],
+                    "networks": [],
                 }
 
                 # Store the previous data to detect deletions and protect against API failures
@@ -418,7 +423,7 @@ async def _async_update_data(self) -> Dict[str, List[Any]]:
                             LOGGER.warning("Authentication failure detected during initial fetch: %s", error_msg)
                             # Trigger auth recovery but continue trying other endpoints
                             if hasattr(self.api, "handle_auth_failure"):
-                                recovery_task = asyncio.create_task(self.api.handle_auth_failure(error_msg))
+                                asyncio.create_task(self.api.handle_auth_failure(error_msg))
 
                             # Preserve previous port forwards data if available
                             if previous_data and "port_forwards" in previous_data and previous_data["port_forwards"]:
@@ -449,6 +454,10 @@ async def _async_update_data(self) -> Dict[str, List[Any]]:
                 await self._update_traffic_rules_in_dict(rules_data)
                 await asyncio.sleep(api_call_delay)
 
+                # Then networks (source of truth for VPN derivation)
+                await self._update_networks_in_dict(rules_data)
+                await asyncio.sleep(api_call_delay)
+
                 # Then legacy firewall rules
                 await self._update_legacy_firewall_rules_in_dict(rules_data)
                 await asyncio.sleep(api_call_delay)
@@ -457,16 +466,18 @@ async def _async_update_data(self) -> Dict[str, List[Any]]:
                 await self._update_qos_rules_in_dict(rules_data)
                 await asyncio.sleep(api_call_delay)
 
-                # Then VPN clients
+                # Then derive VPN clients and servers from networks
                 await self._update_vpn_clients_in_dict(rules_data)
                 await asyncio.sleep(api_call_delay)
-
-                # Then VPN servers
                 await self._update_vpn_servers_in_dict(rules_data)
                 await asyncio.sleep(api_call_delay)
 
                 # Then devices (for LED switches)
                 await self._update_devices_in_dict(rules_data)
+                await asyncio.sleep(api_call_delay)
+
+                # Then port profiles
+                await self._update_port_profiles_in_dict(rules_data)
 
                 # Verify the data is valid - check if we have at least some data in key categories
                 # This helps prevent entity removal during temporary API errors
@@ -582,6 +593,8 @@ async def _async_update_data(self) -> Dict[str, List[Any]]:
                     self.vpn_clients = rules_data.get("vpn_clients", [])
                     self.vpn_servers = rules_data.get("vpn_servers", [])
                     self.devices = rules_data.get("devices", [])
+                    self.port_profiles = rules_data.get("port_profiles", [])
+                    self.networks = rules_data.get("networks", [])
 
                     LOGGER.info("Rule collections after refresh: Port Forwards=%d, Traffic Routes=%d, Firewall Policies=%d, Traffic Rules=%d, Legacy Firewall Rules=%d, WLANs=%d, QoS Rules=%d, VPN Clients=%d, VPN Servers=%d, Devices=%d", 
                                len(self.port_forwards),
@@ -601,10 +614,11 @@ async def _async_update_data(self) -> Dict[str, List[Any]]:
                 LOGGER.error("Error updating coordinator data: %s", err)
 
                 # Check if this is an authentication error
-                auth_error = False
+                # Track auth errors for control flow only (no use afterwards)
+                _auth_error = False
                 error_str = str(err).lower()
                 if "401 unauthorized" in error_str or "403 forbidden" in error_str:
-                    auth_error = True
+                    _auth_error = True
                     self._auth_failures += 1
                     self._authentication_in_progress = True
                     try:
@@ -659,6 +673,8 @@ def _check_for_deleted_rules(self, new_data: Dict[str, List[Any]]) -> None:
             "wlans",
             "vpn_clients",
             "vpn_servers",
+            "port_profiles",
+            "networks",
         ]
 
         for rule_type in all_rule_sources_types:
@@ -806,12 +822,57 @@ async def _update_qos_rules_in_dict(self, data: Dict[str, List[Any]]) -> None:
         await self._update_rule_type_in_dict(data, "qos_rules", self.api.get_qos_rules)
 
     async def _update_vpn_clients_in_dict(self, data: Dict[str, List[Any]]) -> None:
-        """Update VPN clients in the given data dictionary."""
-        await self._update_rule_type_in_dict(data, "vpn_clients", self.api.get_vpn_clients)
+        """Update VPN clients from already-fetched networks."""
+        try:
+            # Derive from networks if available; otherwise, fall back to API
+            networks = data.get("networks") or self.networks
+            if networks:
+                from .models.vpn_config import VPNConfig
+                clients = []
+                for n in networks:
+                    raw = getattr(n, 'raw', {}) if hasattr(n, 'raw') else {}
+                    purpose = raw.get("purpose", "")
+                    vpn_type = raw.get("vpn_type", "")
+                    is_client = purpose == "vpn-client" or vpn_type in ["openvpn-client", "wireguard-client"]
+                    if is_client:
+                        try:
+                            clients.append(VPNConfig(raw))
+                        except Exception as err:
+                            LOGGER.debug("Skipping VPN client conversion error: %s", err)
+                data["vpn_clients"] = clients
+                return
+            # Fallback to API method if networks missing
+            await self._update_rule_type_in_dict(data, "vpn_clients", self.api.get_vpn_clients)
+        except Exception as err:
+            LOGGER.error("Error deriving VPN clients: %s", err)
+            if "vpn_clients" not in data:
+                data["vpn_clients"] = []
 
     async def _update_vpn_servers_in_dict(self, data: Dict[str, List[Any]]) -> None:
-        """Update VPN servers in the given data dictionary."""
-        await self._update_rule_type_in_dict(data, "vpn_servers", self.api.get_vpn_servers)
+        """Update VPN servers from already-fetched networks."""
+        try:
+            networks = data.get("networks") or self.networks
+            if networks:
+                from .models.vpn_config import VPNConfig
+                servers = []
+                for n in networks:
+                    raw = getattr(n, 'raw', {}) if hasattr(n, 'raw') else {}
+                    purpose = raw.get("purpose", "")
+                    vpn_type = raw.get("vpn_type", "")
+                    is_server = purpose == "vpn-server" or vpn_type in ["openvpn-server", "wireguard-server"]
+                    if is_server:
+                        try:
+                            servers.append(VPNConfig(raw))
+                        except Exception as err:
+                            LOGGER.debug("Skipping VPN server conversion error: %s", err)
+                data["vpn_servers"] = servers
+                return
+            # Fallback to API method if networks missing
+            await self._update_rule_type_in_dict(data, "vpn_servers", self.api.get_vpn_servers)
+        except Exception as err:
+            LOGGER.error("Error deriving VPN servers: %s", err)
+            if "vpn_servers" not in data:
+                data["vpn_servers"] = []
 
     async def _update_devices_in_dict(self, data: Dict[str, List[Any]]) -> None:
         """Update devices in the data dictionary."""
@@ -845,6 +906,49 @@ async def _update_devices_in_dict(self, data: Dict[str, List[Any]]) -> None:
             if not hasattr(self, 'devices'):
                 self.devices = []
 
+    async def _update_port_profiles_in_dict(self, data: Dict[str, List[Any]]) -> None:
+        """Update port profiles in the data dictionary and convert to typed objects."""
+        try:
+            LOGGER.info("Fetching port profiles...")
+            future = await self.api.queue_api_operation(self.api.get_port_profiles)
+            profiles = await future if hasattr(future, "__await__") else future
+            typed: List[PortProfile] = []
+            for item in profiles or []:
+                try:
+                    typed.append(PortProfile(item))
+                except Exception as err:
+                    LOGGER.warning("Error converting port profile: %s", err)
+            data["port_profiles"] = typed
+            self.port_profiles = typed
+            LOGGER.info("Updated %d port profiles", len(typed))
+        except Exception as err:
+            LOGGER.error("Failed to update port profiles: %s", err)
+            data["port_profiles"] = []
+            if not hasattr(self, 'port_profiles'):
+                self.port_profiles = []
+
+    async def _update_networks_in_dict(self, data: Dict[str, List[Any]]) -> None:
+        """Update networks list in the data dictionary and convert to typed objects."""
+        try:
+            LOGGER.info("Fetching networks...")
+            future = await self.api.queue_api_operation(self.api.get_networks)
+            networks = await future if hasattr(future, "__await__") else future
+            typed: List[NetworkConf] = []
+            for item in networks or []:
+                try:
+                    # item is already NetworkConf from API layer
+                    typed.append(item)
+                except Exception as err:
+                    LOGGER.warning("Error converting network: %s", err)
+            data["networks"] = typed
+            self.networks = typed
+            LOGGER.info("Updated %d networks", len(typed))
+        except Exception as err:
+            LOGGER.error("Failed to update networks: %s", err)
+            data["networks"] = []
+            if not hasattr(self, 'networks'):
+                self.networks = []
+
     async def _update_rule_type(self, rule_type: str, fetch_method: Callable) -> None:
         """Update a specific rule type in self.data.
 
@@ -994,7 +1098,7 @@ async def _force_refresh_with_cache_clear(self) -> None:
             log_websocket("Starting forced refresh after rule change detected")
 
             # Store previous data for deletion detection
-            previous_data = self.data.copy() if self.data else {}
+            _unused_previous = self.data.copy() if self.data else {}
 
             # Clear the API cache to ensure we get fresh data
             # But do so without disrupting authentication