Skip to content

Add discovery and improve consistency #33

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Jul 28, 2025
Merged

Add discovery and improve consistency #33

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
ced3563
Align logging
CoMPaTech Jul 28, 2025
0581d79
Add discovery class
CoMPaTech Jul 28, 2025
d227e1a
Add discovery class
CoMPaTech Jul 28, 2025
43e098c
Lower coverage requirement, CRAI issues
CoMPaTech Jul 28, 2025
6867740
Tests and improvements
CoMPaTech Jul 28, 2025
029b18e
CRAI suggestions
CoMPaTech Jul 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 16 additions & 16 deletions airos/airos8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
KeyDataMissingError,
)

logger = logging.getLogger(__name__)
_LOGGER = logging.getLogger(__name__)


class AirOS:
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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
Expand All @@ -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 ---
Expand All @@ -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}
Expand All @@ -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
275 changes: 275 additions & 0 deletions airos/discovery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,275 @@
"""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 (AirosEndpointError, AirosListenerError) as err:
# These are expected types of malformed packets. Log the specific error
# and then re-raise as AirosDiscoveryError.
log = f"Parsing failed for packet from {host_ip}: {err}"
_LOGGER.exception(log)
raise AirosDiscoveryError(f"Malformed packet from {host_ip}") from err
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)
raise AirosEndpointError(f"Malformed packet: {log}")

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)
raise AirosEndpointError(f"Malformed packet: {log}")

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)
log = f"Malformed packet: {log}"
raise AirosEndpointError(log)

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)
log = f"Malformed packet: {log}"
raise AirosEndpointError(log)

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)
log = f"Malformed packet: {log}"
raise AirosEndpointError(log)

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}. "
log += f"Cannot determine length, stopping parsing. Remaining: {data[offset - 1 :].hex()}"
_LOGGER.warning(log)
log = f"Malformed packet: {log}"
raise AirosEndpointError(log)

except (struct.error, IndexError) as err:
log = f"Parsing error (struct/index) in AirosDiscoveryProtocol: {err} at offset {offset}. Remaining data: {data[offset:].hex()}"
_LOGGER.debug(log)
log = f"Malformed packet: {log}"
raise AirosEndpointError(log) from err
except AirosEndpointError: # Catch AirosEndpointError specifically, re-raise it
raise
except Exception as err:
_LOGGER.exception("Unexpected error during Airos packet parsing")
raise AirosListenerError from err

return parsed_info
Loading