Add discovery and improve consistency

airos/airos8.py

@@ -18,7 +18,7 @@
     KeyDataMissingError,
 )
 
-logger = logging.getLogger(__name__)
+_LOGGER = logging.getLogger(__name__)
 
 
 class AirOS:
@@ -101,10 +101,10 @@ async def login(self) -> bool:
                 headers=login_request_headers,
             ) as response:
                 if response.status == 403:
-                    logger.error("Authentication denied.")
+                    _LOGGER.error("Authentication denied.")
                     raise ConnectionAuthenticationError from None
                 if not response.cookies:
-                    logger.exception("Empty cookies after login, bailing out.")
+                    _LOGGER.exception("Empty cookies after login, bailing out.")
                     raise ConnectionSetupError from None
                 else:
                     for _, morsel in response.cookies.items():
@@ -155,7 +155,7 @@ async def login(self) -> bool:
                 airos_cookie_found = False
                 ok_cookie_found = False
                 if not self.session.cookie_jar:  # pragma: no cover
-                    logger.exception(
+                    _LOGGER.exception(
                         "COOKIE JAR IS EMPTY after login POST. This is a major issue."
                     )
                     raise ConnectionSetupError from None
@@ -176,24 +176,24 @@ async def login(self) -> bool:
                         self.connected = True
                         return True
                     except json.JSONDecodeError as err:
-                        logger.exception("JSON Decode Error")
+                        _LOGGER.exception("JSON Decode Error")
                         raise DataMissingError from err
 
                 else:
                     log = f"Login failed with status {response.status}. Full Response: {response.text}"
-                    logger.error(log)
+                    _LOGGER.error(log)
                     raise ConnectionAuthenticationError from None
         except (
             aiohttp.ClientError,
             aiohttp.client_exceptions.ConnectionTimeoutError,
         ) as err:
-            logger.exception("Error during login")
+            _LOGGER.exception("Error during login")
             raise DeviceConnectionError from err
 
     async def status(self) -> AirOSData:
         """Retrieve status from the device."""
         if not self.connected:
-            logger.error("Not connected, login first")
+            _LOGGER.error("Not connected, login first")
             raise DeviceConnectionError from None
 
         # --- Step 2: Verify authenticated access by fetching status.cgi ---
@@ -213,32 +213,32 @@ async def status(self) -> AirOSData:
                         try:
                             airos_data = AirOSData.from_dict(response_json)
                         except (MissingField, InvalidFieldValue) as err:
-                            logger.exception("Failed to deserialize AirOS data")
+                            _LOGGER.exception("Failed to deserialize AirOS data")
                             raise KeyDataMissingError from err
 
                         return airos_data
                     except json.JSONDecodeError:
-                        logger.exception(
+                        _LOGGER.exception(
                             "JSON Decode Error in authenticated status response"
                         )
                         raise DataMissingError from None
                 else:
                     log = f"Authenticated status.cgi failed: {response.status}. Response: {response_text}"
-                    logger.error(log)
+                    _LOGGER.error(log)
         except (
             aiohttp.ClientError,
             aiohttp.client_exceptions.ConnectionTimeoutError,
         ) as err:
-            logger.exception("Error during authenticated status.cgi call")
+            _LOGGER.exception("Error during authenticated status.cgi call")
             raise DeviceConnectionError from err
 
     async def stakick(self, mac_address: str = None) -> bool:
         """Reconnect client station."""
         if not self.connected:
-            logger.error("Not connected, login first")
+            _LOGGER.error("Not connected, login first")
             raise DeviceConnectionError from None
         if not mac_address:
-            logger.error("Device mac-address missing")
+            _LOGGER.error("Device mac-address missing")
             raise DataMissingError from None
 
         kick_request_headers = {**self._common_headers}
@@ -262,11 +262,11 @@ async def stakick(self, mac_address: str = None) -> bool:
                     return True
                 response_text = await response.text()
                 log = f"Unable to restart connection response status {response.status} with {response_text}"
-                logger.error(log)
+                _LOGGER.error(log)
                 return False
         except (
             aiohttp.ClientError,
             aiohttp.client_exceptions.ConnectionTimeoutError,
         ) as err:
-            logger.exception("Error during reconnect stakick.cgi call")
+            _LOGGER.exception("Error during reconnect stakick.cgi call")
             raise DeviceConnectionError from err

airos/discovery.py

@@ -0,0 +1,263 @@
+"""Discover Ubiquiti UISP airOS device broadcasts."""
+
+import asyncio
+from collections.abc import Callable
+import logging
+import socket
+import struct
+from typing import Any
+
+from exceptions import AirosDiscoveryError, AirosEndpointError, AirosListenerError
+
+_LOGGER = logging.getLogger(__name__)
+
+DISCOVERY_PORT: int = 10002
+BUFFER_SIZE: int = 1024
+
+
+class AirosDiscoveryProtocol(asyncio.DatagramProtocol):
+    """A UDP protocol implementation for discovering Ubiquiti airOS devices.
+
+    This class listens for UDP broadcast announcements from airOS devices
+    on a specific port (10002) and parses the proprietary packet format
+    to extract device information. It acts as the low-level listener.
+
+    Attributes:
+        callback: An asynchronous callable that will be invoked with
+                  the parsed device information upon discovery.
+        transport: The UDP transport layer object, set once the connection is made.
+
+    """
+
+    def __init__(self, callback: Callable[[dict[str, Any]], None]) -> None:
+        """Initialize AirosDiscoveryProtocol.
+
+        Args:
+            callback: An asynchronous function to call when a device is discovered.
+                      It should accept a dictionary containing device information.
+
+        """
+        self.callback = callback
+        self.transport: asyncio.DatagramTransport | None = None
+
+    def connection_made(self, transport: asyncio.BaseTransport) -> None:
+        """Set up the UDP socket for broadcasting and reusing the address."""
+        self.transport = transport  # type: ignore[assignment] # transport is DatagramTransport
+        sock: socket.socket = self.transport.get_extra_info("socket")
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        log = f"Airos discovery listener (low-level) started on UDP port {DISCOVERY_PORT}."
+        _LOGGER.debug(log)
+
+    def datagram_received(self, data: bytes, addr: tuple[str, int]) -> None:
+        """Parse the received UDP packet and, if successful, schedules the callback.
+
+        Errors during parsing are logged internally by parse_airos_packet.
+        """
+        host_ip: str = addr[0]
+        try:
+            parsed_data: dict[str, Any] | None = self.parse_airos_packet(data, host_ip)
+            if parsed_data:
+                # Schedule the user-provided callback, don't await to keep listener responsive
+                asyncio.create_task(self.callback(parsed_data))  # noqa: RUF006
+        except Exception as err:
+            # General error during datagram reception (e.g., in callback itself)
+            log = f"Error processing Airos discovery packet from {host_ip}. Data hex: {data.hex()}"
+            _LOGGER.exception(log)
+            raise AirosDiscoveryError from err
+
+    def error_received(self, exc: Exception | None) -> None:
+        """Handle send or receive operation raises an OSError."""
+        if exc:
+            log = f"UDP error received in AirosDiscoveryProtocol: {exc}"
+            _LOGGER.error(log)
+
+    def connection_lost(self, exc: Exception | None) -> None:
+        """Handle connection is lost or closed."""
+        _LOGGER.debug("AirosDiscoveryProtocol connection lost.")
+        if exc:
+            _LOGGER.exception("AirosDiscoveryProtocol connection lost due to")
+            raise AirosDiscoveryError from None
+
+    def parse_airos_packet(self, data: bytes, host_ip: str) -> dict[str, Any] | None:
+        """Parse a raw airOS discovery UDP packet.
+
+        This method extracts various pieces of information from the proprietary
+        Ubiquiti airOS discovery packet format, which includes a fixed header
+        followed by a series of Type-Length-Value (TLV) entries. Different
+        TLV types use different length encoding schemes (fixed, 1-byte, 2-byte).
+
+        Args:
+            data: The raw byte data of the UDP packet payload.
+            host_ip: The IP address of the sender, used as a fallback or initial IP.
+
+        Returns:
+            A dictionary containing parsed device information if successful,
+            otherwise None. Values will be None if not found or cannot be parsed.
+
+        """
+        parsed_info: dict[str, str | int | None] = {
+            "ip_address": host_ip,
+            "mac_address": None,
+            "hostname": None,
+            "model": None,
+            "firmware_version": None,
+            "uptime_seconds": None,
+            "ssid": None,
+            "full_model_name": None,
+        }
+
+        # --- Fixed Header (6 bytes) ---
+        if len(data) < 6:
+            log = f"Packet too short for initial fixed header. Length: {len(data)}. Data: {data.hex()}"
+            _LOGGER.debug(log)
+            return None
+
+        if data[0] != 0x01 or data[1] != 0x06:
+            log = f"Packet does not start with expected Airos header (0x01 0x06). Actual: {data[0:2].hex()}"
+            _LOGGER.debug(log)
+            return None
+
+        offset: int = 6
+
+        # --- Main TLV Parsing Loop ---
+        try:
+            while offset < len(data):
+                if (len(data) - offset) < 1:
+                    log = f"Not enough bytes for next TLV type. Remaining: {data[offset:].hex()}"
+                    _LOGGER.debug(log)
+                    break
+
+                tlv_type: int = data[offset]
+                offset += 1
+
+                if tlv_type == 0x06:  # Device MAC Address (fixed 6-byte value)
+                    expected_length: int = 6
+                    if (len(data) - offset) >= expected_length:
+                        mac_bytes: bytes = data[offset : offset + expected_length]
+                        parsed_info["mac_address"] = ":".join(
+                            f"{b:02x}" for b in mac_bytes
+                        ).upper()
+                        offset += expected_length
+                        log = f"Parsed MAC from type 0x06: {parsed_info['mac_address']}"
+                        _LOGGER.debug(log)
+                    else:
+                        log = f"Truncated MAC address TLV (Type 0x06). Expected {expected_length}, got {len(data) - offset} bytes. Remaining: {data[offset:].hex()}"
+                        _LOGGER.warning(log)
+                        break
+
+                elif tlv_type in [
+                    0x02,
+                    0x03,
+                    0x0A,
+                    0x0B,
+                    0x0C,
+                    0x0D,
+                    0x0E,
+                    0x10,
+                    0x14,
+                    0x18,
+                ]:
+                    if (len(data) - offset) < 2:
+                        log = f"Truncated TLV (Type {tlv_type:#x}), no 2-byte length field. Remaining: {data[offset:].hex()}"
+                        _LOGGER.warning(log)
+                        break
+
+                    tlv_length: int = struct.unpack_from(">H", data, offset)[0]
+                    offset += 2
+
+                    if tlv_length > (len(data) - offset):
+                        log = f"TLV type {tlv_type:#x} length {tlv_length} exceeds remaining data "
+                        _LOGGER.warning(log)
+                        log = f"({len(data) - offset} bytes left). Packet malformed. "
+                        _LOGGER.warning(log)
+                        log = f"Data from TLV start: {data[offset - 3 :].hex()}"
+                        _LOGGER.warning(log)
+                        break
+
+                    tlv_value: bytes = data[offset : offset + tlv_length]
+
+                    if tlv_type == 0x02:
+                        if tlv_length == 10:
+                            ip_bytes: bytes = tlv_value[6:10]
+                            parsed_info["ip_address"] = ".".join(map(str, ip_bytes))
+                            log = f"Parsed IP from type 0x02 block: {parsed_info['ip_address']}"
+                            _LOGGER.debug(log)
+                        else:
+                            log = f"Unexpected length for 0x02 TLV (MAC+IP). Expected 10, got {tlv_length}. Value: {tlv_value.hex()}"
+                            _LOGGER.warning(log)
+
+                    elif tlv_type == 0x03:
+                        parsed_info["firmware_version"] = tlv_value.decode(
+                            "ascii", errors="ignore"
+                        )
+                        log = f"Parsed Firmware: {parsed_info['firmware_version']}"
+                        _LOGGER.debug(log)
+
+                    elif tlv_type == 0x0A:
+                        if tlv_length == 4:
+                            parsed_info["uptime_seconds"] = struct.unpack(
+                                ">I", tlv_value
+                            )[0]
+                            log = f"Parsed Uptime: {parsed_info['uptime_seconds']}s"
+                            _LOGGER.debug(log)
+                        else:
+                            log = f"Unexpected length for Uptime (Type 0x0A): {tlv_length}. Value: {tlv_value.hex()}"
+                            _LOGGER.warning(log)
+
+                    elif tlv_type == 0x0B:
+                        parsed_info["hostname"] = tlv_value.decode(
+                            "utf-8", errors="ignore"
+                        )
+                        log = f"Parsed Hostname: {parsed_info['hostname']}"
+                        _LOGGER.debug(log)
+
+                    elif tlv_type == 0x0C:
+                        parsed_info["model"] = tlv_value.decode(
+                            "ascii", errors="ignore"
+                        )
+                        log = f"Parsed Model: {parsed_info['model']}"
+                        _LOGGER.debug(log)
+
+                    elif tlv_type == 0x0D:
+                        parsed_info["ssid"] = tlv_value.decode("utf-8", errors="ignore")
+                        log = f"Parsed SSID: {parsed_info['ssid']}"
+                        _LOGGER.debug(log)
+
+                    elif tlv_type == 0x14:
+                        parsed_info["full_model_name"] = tlv_value.decode(
+                            "utf-8", errors="ignore"
+                        )
+                        log = (
+                            f"Parsed Full Model Name: {parsed_info['full_model_name']}"
+                        )
+                        _LOGGER.debug(log)
+
+                    elif tlv_type == 0x18:
+                        if tlv_length == 4 and tlv_value == b"\x00\x00\x00\x00":
+                            _LOGGER.debug("Detected end marker (Type 0x18).")
+                        else:
+                            log = f"Unhandled TLV type: {tlv_type:#x} with length {tlv_length}. Value: {tlv_value.hex()}"
+                            _LOGGER.debug(log)
+                    elif tlv_type in [0x0E, 0x10]:
+                        log = f"Unhandled TLV type: {tlv_type:#x} with length {tlv_length}. Value: {tlv_value.hex()}"
+                        _LOGGER.debug(log)
+
+                    offset += tlv_length
+
+                else:
+                    log = f"Unhandled TLV type: {tlv_type:#x} at offset {offset - 1}. "
+                    _LOGGER.warning(log)
+                    log = f"Cannot determine length, stopping parsing. Remaining: {data[offset - 1 :].hex()}"
+                    _LOGGER.warning(log)
+                    break
+
+        except (struct.error, IndexError) as err:
+            log = f"Parsing error (struct/index) in AirosDiscoveryProtocol: {err} at offset {offset}. Remaining data: {data[offset:].hex()}"
+            _LOGGER.warning(log)
+            raise AirosEndpointError from err
+        except Exception as err:
+            _LOGGER.exception("Unexpected error during Airos packet parsing")
+            raise AirosListenerError from err
+
+        return parsed_info

airos/exceptions.py

@@ -23,3 +23,15 @@ class KeyDataMissingError(AirOSException):
 
 class DeviceConnectionError(AirOSException):
     """Raised when unable to connect."""
+
+
+class AirosDiscoveryError(AirOSException):
+    """Base exception for Airos discovery issues."""
+
+
+class AirosListenerError(AirosDiscoveryError):
+    """Raised when the Airos listener encounters an error."""
+
+
+class AirosEndpointError(AirosDiscoveryError):
+    """Raised when there's an issue with the network endpoint."""