Skip to content

Commit 0581d79

Browse files
CoMPaTechCoMPaTech
committed
Add discovery class
1 parent ced3563 commit 0581d79

File tree

2 files changed

+264
-1
lines changed

2 files changed

+264
-1
lines changed

‎airos/discovery.py‎

Lines changed: 263 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,263 @@
1+
"""Discover Ubiquiti UISP airOS device broadcasts."""
2+
3+
import asyncio
4+
from collections.abc import Callable
5+
import logging
6+
import socket
7+
import struct
8+
from typing import Any
9+
10+
from exceptions import AirosDiscoveryError, AirosEndpointError, AirosListenerError
11+
12+
_LOGGER = logging.getLogger(__name__)
13+
14+
DISCOVERY_PORT: int = 10002
15+
BUFFER_SIZE: int = 1024
16+
17+
18+
class AirosDiscoveryProtocol(asyncio.DatagramProtocol):
19+
"""A UDP protocol implementation for discovering Ubiquiti airOS devices.
20+
21+
This class listens for UDP broadcast announcements from airOS devices
22+
on a specific port (10002) and parses the proprietary packet format
23+
to extract device information. It acts as the low-level listener.
24+
25+
Attributes:
26+
callback: An asynchronous callable that will be invoked with
27+
the parsed device information upon discovery.
28+
transport: The UDP transport layer object, set once the connection is made.
29+
30+
"""
31+
32+
def __init__(self, callback: Callable[[dict[str, Any]], None]) -> None:
33+
"""Initialize AirosDiscoveryProtocol.
34+
35+
Args:
36+
callback: An asynchronous function to call when a device is discovered.
37+
It should accept a dictionary containing device information.
38+
39+
"""
40+
self.callback = callback
41+
self.transport: asyncio.DatagramTransport | None = None
42+
43+
def connection_made(self, transport: asyncio.BaseTransport) -> None:
44+
"""Set up the UDP socket for broadcasting and reusing the address."""
45+
self.transport = transport # type: ignore[assignment] # transport is DatagramTransport
46+
sock: socket.socket = self.transport.get_extra_info("socket")
47+
sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
48+
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
49+
log = f"Airos discovery listener (low-level) started on UDP port {DISCOVERY_PORT}."
50+
_LOGGER.debug(log)
51+
52+
def datagram_received(self, data: bytes, addr: tuple[str, int]) -> None:
53+
"""Parse the received UDP packet and, if successful, schedules the callback.
54+
55+
Errors during parsing are logged internally by parse_airos_packet.
56+
"""
57+
host_ip: str = addr[0]
58+
try:
59+
parsed_data: dict[str, Any] | None = self.parse_airos_packet(data, host_ip)
60+
if parsed_data:
61+
# Schedule the user-provided callback, don't await to keep listener responsive
62+
asyncio.create_task(self.callback(parsed_data)) # noqa: RUF006
63+
except Exception as err:
64+
# General error during datagram reception (e.g., in callback itself)
65+
log = f"Error processing Airos discovery packet from {host_ip}. Data hex: {data.hex()}"
66+
_LOGGER.exception(log)
67+
raise AirosDiscoveryError from err
68+
69+
def error_received(self, exc: Exception | None) -> None:
70+
"""Handle send or receive operation raises an OSError."""
71+
if exc:
72+
log = f"UDP error received in AirosDiscoveryProtocol: {exc}"
73+
_LOGGER.error(log)
74+
75+
def connection_lost(self, exc: Exception | None) -> None:
76+
"""Handle connection is lost or closed."""
77+
_LOGGER.debug("AirosDiscoveryProtocol connection lost.")
78+
if exc:
79+
_LOGGER.exception("AirosDiscoveryProtocol connection lost due to")
80+
raise AirosDiscoveryError from None
81+
82+
def parse_airos_packet(self, data: bytes, host_ip: str) -> dict[str, Any] | None:
83+
"""Parse a raw airOS discovery UDP packet.
84+
85+
This method extracts various pieces of information from the proprietary
86+
Ubiquiti airOS discovery packet format, which includes a fixed header
87+
followed by a series of Type-Length-Value (TLV) entries. Different
88+
TLV types use different length encoding schemes (fixed, 1-byte, 2-byte).
89+
90+
Args:
91+
data: The raw byte data of the UDP packet payload.
92+
host_ip: The IP address of the sender, used as a fallback or initial IP.
93+
94+
Returns:
95+
A dictionary containing parsed device information if successful,
96+
otherwise None. Values will be None if not found or cannot be parsed.
97+
98+
"""
99+
parsed_info: dict[str, str | int | None] = {
100+
"ip_address": host_ip,
101+
"mac_address": None,
102+
"hostname": None,
103+
"model": None,
104+
"firmware_version": None,
105+
"uptime_seconds": None,
106+
"ssid": None,
107+
"full_model_name": None,
108+
}
109+
110+
# --- Fixed Header (6 bytes) ---
111+
if len(data) < 6:
112+
log = f"Packet too short for initial fixed header. Length: {len(data)}. Data: {data.hex()}"
113+
_LOGGER.debug(log)
114+
return None
115+
116+
if data[0] != 0x01 or data[1] != 0x06:
117+
log = f"Packet does not start with expected Airos header (0x01 0x06). Actual: {data[0:2].hex()}"
118+
_LOGGER.debug(log)
119+
return None
120+
121+
offset: int = 6
122+
123+
# --- Main TLV Parsing Loop ---
124+
try:
125+
while offset < len(data):
126+
if (len(data) - offset) < 1:
127+
log = f"Not enough bytes for next TLV type. Remaining: {data[offset:].hex()}"
128+
_LOGGER.debug(log)
129+
break
130+
131+
tlv_type: int = data[offset]
132+
offset += 1
133+
134+
if tlv_type == 0x06: # Device MAC Address (fixed 6-byte value)
135+
expected_length: int = 6
136+
if (len(data) - offset) >= expected_length:
137+
mac_bytes: bytes = data[offset : offset + expected_length]
138+
parsed_info["mac_address"] = ":".join(
139+
f"{b:02x}" for b in mac_bytes
140+
).upper()
141+
offset += expected_length
142+
log = f"Parsed MAC from type 0x06: {parsed_info['mac_address']}"
143+
_LOGGER.debug(log)
144+
else:
145+
log = f"Truncated MAC address TLV (Type 0x06). Expected {expected_length}, got {len(data) - offset} bytes. Remaining: {data[offset:].hex()}"
146+
_LOGGER.warning(log)
147+
break
148+
149+
elif tlv_type in [
150+
0x02,
151+
0x03,
152+
0x0A,
153+
0x0B,
154+
0x0C,
155+
0x0D,
156+
0x0E,
157+
0x10,
158+
0x14,
159+
0x18,
160+
]:
161+
if (len(data) - offset) < 2:
162+
log = f"Truncated TLV (Type {tlv_type:#x}), no 2-byte length field. Remaining: {data[offset:].hex()}"
163+
_LOGGER.warning(log)
164+
break
165+
166+
tlv_length: int = struct.unpack_from(">H", data, offset)[0]
167+
offset += 2
168+
169+
if tlv_length > (len(data) - offset):
170+
log = f"TLV type {tlv_type:#x} length {tlv_length} exceeds remaining data "
171+
_LOGGER.warning(log)
172+
log = f"({len(data) - offset} bytes left). Packet malformed. "
173+
_LOGGER.warning(log)
174+
log = f"Data from TLV start: {data[offset - 3 :].hex()}"
175+
_LOGGER.warning(log)
176+
break
177+
178+
tlv_value: bytes = data[offset : offset + tlv_length]
179+
180+
if tlv_type == 0x02:
181+
if tlv_length == 10:
182+
ip_bytes: bytes = tlv_value[6:10]
183+
parsed_info["ip_address"] = ".".join(map(str, ip_bytes))
184+
log = f"Parsed IP from type 0x02 block: {parsed_info['ip_address']}"
185+
_LOGGER.debug(log)
186+
else:
187+
log = f"Unexpected length for 0x02 TLV (MAC+IP). Expected 10, got {tlv_length}. Value: {tlv_value.hex()}"
188+
_LOGGER.warning(log)
189+
190+
elif tlv_type == 0x03:
191+
parsed_info["firmware_version"] = tlv_value.decode(
192+
"ascii", errors="ignore"
193+
)
194+
log = f"Parsed Firmware: {parsed_info['firmware_version']}"
195+
_LOGGER.debug(log)
196+
197+
elif tlv_type == 0x0A:
198+
if tlv_length == 4:
199+
parsed_info["uptime_seconds"] = struct.unpack(
200+
">I", tlv_value
201+
)[0]
202+
log = f"Parsed Uptime: {parsed_info['uptime_seconds']}s"
203+
_LOGGER.debug(log)
204+
else:
205+
log = f"Unexpected length for Uptime (Type 0x0A): {tlv_length}. Value: {tlv_value.hex()}"
206+
_LOGGER.warning(log)
207+
208+
elif tlv_type == 0x0B:
209+
parsed_info["hostname"] = tlv_value.decode(
210+
"utf-8", errors="ignore"
211+
)
212+
log = f"Parsed Hostname: {parsed_info['hostname']}"
213+
_LOGGER.debug(log)
214+
215+
elif tlv_type == 0x0C:
216+
parsed_info["model"] = tlv_value.decode(
217+
"ascii", errors="ignore"
218+
)
219+
log = f"Parsed Model: {parsed_info['model']}"
220+
_LOGGER.debug(log)
221+
222+
elif tlv_type == 0x0D:
223+
parsed_info["ssid"] = tlv_value.decode("utf-8", errors="ignore")
224+
log = f"Parsed SSID: {parsed_info['ssid']}"
225+
_LOGGER.debug(log)
226+
227+
elif tlv_type == 0x14:
228+
parsed_info["full_model_name"] = tlv_value.decode(
229+
"utf-8", errors="ignore"
230+
)
231+
log = (
232+
f"Parsed Full Model Name: {parsed_info['full_model_name']}"
233+
)
234+
_LOGGER.debug(log)
235+
236+
elif tlv_type == 0x18:
237+
if tlv_length == 4 and tlv_value == b"\x00\x00\x00\x00":
238+
_LOGGER.debug("Detected end marker (Type 0x18).")
239+
else:
240+
log = f"Unhandled TLV type: {tlv_type:#x} with length {tlv_length}. Value: {tlv_value.hex()}"
241+
_LOGGER.debug(log)
242+
elif tlv_type in [0x0E, 0x10]:
243+
log = f"Unhandled TLV type: {tlv_type:#x} with length {tlv_length}. Value: {tlv_value.hex()}"
244+
_LOGGER.debug(log)
245+
246+
offset += tlv_length
247+
248+
else:
249+
log = f"Unhandled TLV type: {tlv_type:#x} at offset {offset - 1}. "
250+
_LOGGER.warning(log)
251+
log = f"Cannot determine length, stopping parsing. Remaining: {data[offset - 1 :].hex()}"
252+
_LOGGER.warning(log)
253+
break
254+
255+
except (struct.error, IndexError) as err:
256+
log = f"Parsing error (struct/index) in AirosDiscoveryProtocol: {err} at offset {offset}. Remaining data: {data[offset:].hex()}"
257+
_LOGGER.warning(log)
258+
raise AirosEndpointError from err
259+
except Exception as err:
260+
_LOGGER.exception("Unexpected error during Airos packet parsing")
261+
raise AirosListenerError from err
262+
263+
return parsed_info

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
44

55
[project]
66
name = "airos"
7-
version = "0.1.8"
7+
version = "0.2.0"
88
license = "MIT"
99
description = "Ubiquity airOS module(s) for Python 3."
1010
readme = "README.md"

0 commit comments

Comments
 (0)