Skip to content

Commit a01ac76

Browse files
CoMPaTechCoMPaTech
authored
Merge pull request #33 from CoMPaTech/testing
Add discovery and improve consistency
2 parents 14b684b + 029b18e commit a01ac76

File tree

7 files changed

+614
-17
lines changed

7 files changed

+614
-17
lines changed

‎airos/airos8.py‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
KeyDataMissingError,
1919
)
2020

21-
logger = logging.getLogger(__name__)
21+
_LOGGER = logging.getLogger(__name__)
2222

2323

2424
class AirOS:
@@ -101,10 +101,10 @@ async def login(self) -> bool:
101101
headers=login_request_headers,
102102
) as response:
103103
if response.status == 403:
104-
logger.error("Authentication denied.")
104+
_LOGGER.error("Authentication denied.")
105105
raise ConnectionAuthenticationError from None
106106
if not response.cookies:
107-
logger.exception("Empty cookies after login, bailing out.")
107+
_LOGGER.exception("Empty cookies after login, bailing out.")
108108
raise ConnectionSetupError from None
109109
else:
110110
for _, morsel in response.cookies.items():
@@ -155,7 +155,7 @@ async def login(self) -> bool:
155155
airos_cookie_found = False
156156
ok_cookie_found = False
157157
if not self.session.cookie_jar: # pragma: no cover
158-
logger.exception(
158+
_LOGGER.exception(
159159
"COOKIE JAR IS EMPTY after login POST. This is a major issue."
160160
)
161161
raise ConnectionSetupError from None
@@ -176,24 +176,24 @@ async def login(self) -> bool:
176176
self.connected = True
177177
return True
178178
except json.JSONDecodeError as err:
179-
logger.exception("JSON Decode Error")
179+
_LOGGER.exception("JSON Decode Error")
180180
raise DataMissingError from err
181181

182182
else:
183183
log = f"Login failed with status {response.status}. Full Response: {response.text}"
184-
logger.error(log)
184+
_LOGGER.error(log)
185185
raise ConnectionAuthenticationError from None
186186
except (
187187
aiohttp.ClientError,
188188
aiohttp.client_exceptions.ConnectionTimeoutError,
189189
) as err:
190-
logger.exception("Error during login")
190+
_LOGGER.exception("Error during login")
191191
raise DeviceConnectionError from err
192192

193193
async def status(self) -> AirOSData:
194194
"""Retrieve status from the device."""
195195
if not self.connected:
196-
logger.error("Not connected, login first")
196+
_LOGGER.error("Not connected, login first")
197197
raise DeviceConnectionError from None
198198

199199
# --- Step 2: Verify authenticated access by fetching status.cgi ---
@@ -213,32 +213,32 @@ async def status(self) -> AirOSData:
213213
try:
214214
airos_data = AirOSData.from_dict(response_json)
215215
except (MissingField, InvalidFieldValue) as err:
216-
logger.exception("Failed to deserialize AirOS data")
216+
_LOGGER.exception("Failed to deserialize AirOS data")
217217
raise KeyDataMissingError from err
218218

219219
return airos_data
220220
except json.JSONDecodeError:
221-
logger.exception(
221+
_LOGGER.exception(
222222
"JSON Decode Error in authenticated status response"
223223
)
224224
raise DataMissingError from None
225225
else:
226226
log = f"Authenticated status.cgi failed: {response.status}. Response: {response_text}"
227-
logger.error(log)
227+
_LOGGER.error(log)
228228
except (
229229
aiohttp.ClientError,
230230
aiohttp.client_exceptions.ConnectionTimeoutError,
231231
) as err:
232-
logger.exception("Error during authenticated status.cgi call")
232+
_LOGGER.exception("Error during authenticated status.cgi call")
233233
raise DeviceConnectionError from err
234234

235235
async def stakick(self, mac_address: str = None) -> bool:
236236
"""Reconnect client station."""
237237
if not self.connected:
238-
logger.error("Not connected, login first")
238+
_LOGGER.error("Not connected, login first")
239239
raise DeviceConnectionError from None
240240
if not mac_address:
241-
logger.error("Device mac-address missing")
241+
_LOGGER.error("Device mac-address missing")
242242
raise DataMissingError from None
243243

244244
kick_request_headers = {**self._common_headers}
@@ -262,11 +262,11 @@ async def stakick(self, mac_address: str = None) -> bool:
262262
return True
263263
response_text = await response.text()
264264
log = f"Unable to restart connection response status {response.status} with {response_text}"
265-
logger.error(log)
265+
_LOGGER.error(log)
266266
return False
267267
except (
268268
aiohttp.ClientError,
269269
aiohttp.client_exceptions.ConnectionTimeoutError,
270270
) as err:
271-
logger.exception("Error during reconnect stakick.cgi call")
271+
_LOGGER.exception("Error during reconnect stakick.cgi call")
272272
raise DeviceConnectionError from err

‎airos/discovery.py‎

Lines changed: 275 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,275 @@
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 (AirosEndpointError, AirosListenerError) as err:
64+
# These are expected types of malformed packets. Log the specific error
65+
# and then re-raise as AirosDiscoveryError.
66+
log = f"Parsing failed for packet from {host_ip}: {err}"
67+
_LOGGER.exception(log)
68+
raise AirosDiscoveryError(f"Malformed packet from {host_ip}") from err
69+
except Exception as err:
70+
# General error during datagram reception (e.g., in callback itself)
71+
log = f"Error processing Airos discovery packet from {host_ip}. Data hex: {data.hex()}"
72+
_LOGGER.exception(log)
73+
raise AirosDiscoveryError from err
74+
75+
def error_received(self, exc: Exception | None) -> None:
76+
"""Handle send or receive operation raises an OSError."""
77+
if exc:
78+
log = f"UDP error received in AirosDiscoveryProtocol: {exc}"
79+
_LOGGER.error(log)
80+
81+
def connection_lost(self, exc: Exception | None) -> None:
82+
"""Handle connection is lost or closed."""
83+
_LOGGER.debug("AirosDiscoveryProtocol connection lost.")
84+
if exc:
85+
_LOGGER.exception("AirosDiscoveryProtocol connection lost due to")
86+
raise AirosDiscoveryError from None
87+
88+
def parse_airos_packet(self, data: bytes, host_ip: str) -> dict[str, Any] | None:
89+
"""Parse a raw airOS discovery UDP packet.
90+
91+
This method extracts various pieces of information from the proprietary
92+
Ubiquiti airOS discovery packet format, which includes a fixed header
93+
followed by a series of Type-Length-Value (TLV) entries. Different
94+
TLV types use different length encoding schemes (fixed, 1-byte, 2-byte).
95+
96+
Args:
97+
data: The raw byte data of the UDP packet payload.
98+
host_ip: The IP address of the sender, used as a fallback or initial IP.
99+
100+
Returns:
101+
A dictionary containing parsed device information if successful,
102+
otherwise None. Values will be None if not found or cannot be parsed.
103+
104+
"""
105+
parsed_info: dict[str, str | int | None] = {
106+
"ip_address": host_ip,
107+
"mac_address": None,
108+
"hostname": None,
109+
"model": None,
110+
"firmware_version": None,
111+
"uptime_seconds": None,
112+
"ssid": None,
113+
"full_model_name": None,
114+
}
115+
116+
# --- Fixed Header (6 bytes) ---
117+
if len(data) < 6:
118+
log = f"Packet too short for initial fixed header. Length: {len(data)}. Data: {data.hex()}"
119+
_LOGGER.debug(log)
120+
raise AirosEndpointError(f"Malformed packet: {log}")
121+
122+
if data[0] != 0x01 or data[1] != 0x06:
123+
log = f"Packet does not start with expected Airos header (0x01 0x06). Actual: {data[0:2].hex()}"
124+
_LOGGER.debug(log)
125+
raise AirosEndpointError(f"Malformed packet: {log}")
126+
127+
offset: int = 6
128+
129+
# --- Main TLV Parsing Loop ---
130+
try:
131+
while offset < len(data):
132+
if (len(data) - offset) < 1:
133+
log = f"Not enough bytes for next TLV type. Remaining: {data[offset:].hex()}"
134+
_LOGGER.debug(log)
135+
break
136+
137+
tlv_type: int = data[offset]
138+
offset += 1
139+
140+
if tlv_type == 0x06: # Device MAC Address (fixed 6-byte value)
141+
expected_length: int = 6
142+
if (len(data) - offset) >= expected_length:
143+
mac_bytes: bytes = data[offset : offset + expected_length]
144+
parsed_info["mac_address"] = ":".join(
145+
f"{b:02x}" for b in mac_bytes
146+
).upper()
147+
offset += expected_length
148+
log = f"Parsed MAC from type 0x06: {parsed_info['mac_address']}"
149+
_LOGGER.debug(log)
150+
else:
151+
log = f"Truncated MAC address TLV (Type 0x06). Expected {expected_length}, got {len(data) - offset} bytes. Remaining: {data[offset:].hex()}"
152+
_LOGGER.warning(log)
153+
log = f"Malformed packet: {log}"
154+
raise AirosEndpointError(log)
155+
156+
elif tlv_type in [
157+
0x02,
158+
0x03,
159+
0x0A,
160+
0x0B,
161+
0x0C,
162+
0x0D,
163+
0x0E,
164+
0x10,
165+
0x14,
166+
0x18,
167+
]:
168+
if (len(data) - offset) < 2:
169+
log = f"Truncated TLV (Type {tlv_type:#x}), no 2-byte length field. Remaining: {data[offset:].hex()}"
170+
_LOGGER.warning(log)
171+
log = f"Malformed packet: {log}"
172+
raise AirosEndpointError(log)
173+
174+
tlv_length: int = struct.unpack_from(">H", data, offset)[0]
175+
offset += 2
176+
177+
if tlv_length > (len(data) - offset):
178+
log = f"TLV type {tlv_type:#x} length {tlv_length} exceeds remaining data "
179+
_LOGGER.warning(log)
180+
log = f"({len(data) - offset} bytes left). Packet malformed. "
181+
_LOGGER.warning(log)
182+
log = f"Data from TLV start: {data[offset - 3 :].hex()}"
183+
_LOGGER.warning(log)
184+
log = f"Malformed packet: {log}"
185+
raise AirosEndpointError(log)
186+
187+
tlv_value: bytes = data[offset : offset + tlv_length]
188+
189+
if tlv_type == 0x02:
190+
if tlv_length == 10:
191+
ip_bytes: bytes = tlv_value[6:10]
192+
parsed_info["ip_address"] = ".".join(map(str, ip_bytes))
193+
log = f"Parsed IP from type 0x02 block: {parsed_info['ip_address']}"
194+
_LOGGER.debug(log)
195+
else:
196+
log = f"Unexpected length for 0x02 TLV (MAC+IP). Expected 10, got {tlv_length}. Value: {tlv_value.hex()}"
197+
_LOGGER.warning(log)
198+
199+
elif tlv_type == 0x03:
200+
parsed_info["firmware_version"] = tlv_value.decode(
201+
"ascii", errors="ignore"
202+
)
203+
log = f"Parsed Firmware: {parsed_info['firmware_version']}"
204+
_LOGGER.debug(log)
205+
206+
elif tlv_type == 0x0A:
207+
if tlv_length == 4:
208+
parsed_info["uptime_seconds"] = struct.unpack(
209+
">I", tlv_value
210+
)[0]
211+
log = f"Parsed Uptime: {parsed_info['uptime_seconds']}s"
212+
_LOGGER.debug(log)
213+
else:
214+
log = f"Unexpected length for Uptime (Type 0x0A): {tlv_length}. Value: {tlv_value.hex()}"
215+
_LOGGER.warning(log)
216+
217+
elif tlv_type == 0x0B:
218+
parsed_info["hostname"] = tlv_value.decode(
219+
"utf-8", errors="ignore"
220+
)
221+
log = f"Parsed Hostname: {parsed_info['hostname']}"
222+
_LOGGER.debug(log)
223+
224+
elif tlv_type == 0x0C:
225+
parsed_info["model"] = tlv_value.decode(
226+
"ascii", errors="ignore"
227+
)
228+
log = f"Parsed Model: {parsed_info['model']}"
229+
_LOGGER.debug(log)
230+
231+
elif tlv_type == 0x0D:
232+
parsed_info["ssid"] = tlv_value.decode("utf-8", errors="ignore")
233+
log = f"Parsed SSID: {parsed_info['ssid']}"
234+
_LOGGER.debug(log)
235+
236+
elif tlv_type == 0x14:
237+
parsed_info["full_model_name"] = tlv_value.decode(
238+
"utf-8", errors="ignore"
239+
)
240+
log = (
241+
f"Parsed Full Model Name: {parsed_info['full_model_name']}"
242+
)
243+
_LOGGER.debug(log)
244+
245+
elif tlv_type == 0x18:
246+
if tlv_length == 4 and tlv_value == b"\x00\x00\x00\x00":
247+
_LOGGER.debug("Detected end marker (Type 0x18).")
248+
else:
249+
log = f"Unhandled TLV type: {tlv_type:#x} with length {tlv_length}. Value: {tlv_value.hex()}"
250+
_LOGGER.debug(log)
251+
elif tlv_type in [0x0E, 0x10]:
252+
log = f"Unhandled TLV type: {tlv_type:#x} with length {tlv_length}. Value: {tlv_value.hex()}"
253+
_LOGGER.debug(log)
254+
255+
offset += tlv_length
256+
257+
else:
258+
log = f"Unhandled TLV type: {tlv_type:#x} at offset {offset - 1}. "
259+
log += f"Cannot determine length, stopping parsing. Remaining: {data[offset - 1 :].hex()}"
260+
_LOGGER.warning(log)
261+
log = f"Malformed packet: {log}"
262+
raise AirosEndpointError(log)
263+
264+
except (struct.error, IndexError) as err:
265+
log = f"Parsing error (struct/index) in AirosDiscoveryProtocol: {err} at offset {offset}. Remaining data: {data[offset:].hex()}"
266+
_LOGGER.debug(log)
267+
log = f"Malformed packet: {log}"
268+
raise AirosEndpointError(log) from err
269+
except AirosEndpointError: # Catch AirosEndpointError specifically, re-raise it
270+
raise
271+
except Exception as err:
272+
_LOGGER.exception("Unexpected error during Airos packet parsing")
273+
raise AirosListenerError from err
274+
275+
return parsed_info

0 commit comments

Comments
 (0)