From 1e862136804f249e97857a67d00e3a0c91c1e5a6 Mon Sep 17 00:00:00 2001 From: Tom Scholten Date: Sat, 19 Jul 2025 20:07:28 +0200 Subject: [PATCH 1/2] Add tests and station reconnect --- .gitignore | 1 + airos/airos8.py | 57 ++++++++++++++++++++++++--- pyproject.toml | 2 +- tests/test_stations.py | 87 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 141 insertions(+), 6 deletions(-) diff --git a/.gitignore b/.gitignore index 94f98cd..8a42694 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ tests/__pycache__ .vscode .coverage tmp +todo diff --git a/airos/airos8.py b/airos/airos8.py index e57b73d..345ae10 100644 --- a/airos/airos8.py +++ b/airos/airos8.py @@ -47,6 +47,7 @@ def __init__( self._login_url = f"{self.base_url}/api/auth" # AirOS 8 self._status_cgi_url = f"{self.base_url}/status.cgi" # AirOS 8 + self._stakick_cgi_url = f"{self.base_url}/stakick.cgi" # AirOS 8 self.current_csrf_token = None self._use_json_for_login_post = False @@ -147,19 +148,19 @@ async def login(self) -> bool: # Re-check cookies in self.session.cookie_jar AFTER potential manual injection airos_cookie_found = False ok_cookie_found = False - if not self.session.cookie_jar: + if not self.session.cookie_jar: # pragma: no cover logger.exception( "COOKIE JAR IS EMPTY after login POST. This is a major issue." ) raise ConnectionSetupError from None - for cookie in self.session.cookie_jar: + for cookie in self.session.cookie_jar: # pragma: no cover if cookie.key.startswith("AIROS_"): airos_cookie_found = True if cookie.key == "ok": ok_cookie_found = True if not airos_cookie_found and not ok_cookie_found: - raise ConnectionSetupError from None + raise ConnectionSetupError from None # pragma: no cover response_text = await response.text() @@ -176,7 +177,10 @@ async def login(self) -> bool: log = f"Login failed with status {response.status}. Full Response: {response.text}" logger.error(log) raise ConnectionAuthenticationError from None - except aiohttp.ClientError as err: + except ( + aiohttp.ClientError, + aiohttp.client_exceptions.ConnectionTimeoutError, + ) as err: logger.exception("Error during login") raise DeviceConnectionError from err @@ -208,6 +212,49 @@ async def status(self) -> dict: else: log = f"Authenticated status.cgi failed: {response.status}. Response: {response_text}" logger.error(log) - except aiohttp.ClientError as err: + except ( + aiohttp.ClientError, + aiohttp.client_exceptions.ConnectionTimeoutError, + ) as err: 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") + raise DeviceConnectionError from None + if not mac_address: + logger.error("Device mac-address missing") + raise DataMissingError from None + + # --- Step 2: Verify authenticated access by fetching status.cgi --- + kick_request_headers = {**self._common_headers} + if self.current_csrf_token: + kick_request_headers["X-CSRF-ID"] = self.current_csrf_token + + kick_payload = { + "staif": "ath0", + "staid": urlparse(mac_address.upper()), + } + + kick_request_headers["Content-Type"] = ( + "application/x-www-form-urlencoded; charset=UTF-8" + ) + post_data = kick_payload + + try: + async with self.session.post( + self._stakick_cgi_url, + headers=kick_request_headers, + data=post_data, + ) as response: + if response.status == 200: + return True + return False + except ( + aiohttp.ClientError, + aiohttp.client_exceptions.ConnectionTimeoutError, + ) as err: + logger.exception("Error during reconnect stakick.cgi call") + raise DeviceConnectionError from err diff --git a/pyproject.toml b/pyproject.toml index 30682c0..699af56 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "airos" -version = "0.0.8" +version = "0.0.9" license = "MIT" description = "Ubiquity airOS module(s) for Python 3." readme = "README.md" diff --git a/tests/test_stations.py b/tests/test_stations.py index daa48e6..7758074 100644 --- a/tests/test_stations.py +++ b/tests/test_stations.py @@ -5,9 +5,11 @@ import os from unittest.mock import AsyncMock, MagicMock, patch +import airos.exceptions import pytest import aiofiles +import aiohttp async def _read_fixture(fixture: str = "ap-ptp"): @@ -56,3 +58,88 @@ async def test_ap(airos_device, base_url, mode): # Verify the fixture returns the correct mode assert status.get("wireless", {}).get("mode") == mode + + +@pytest.mark.asyncio +async def test_reconnect(airos_device, base_url): + """Test reconnect client.""" + # --- Prepare fake POST /api/stakick response --- + mock_stakick_response = MagicMock() + mock_stakick_response.__aenter__.return_value = mock_stakick_response + mock_stakick_response.status = 200 + + with ( + patch.object(airos_device.session, "post", return_value=mock_stakick_response), + patch.object(airos_device, "connected", return_value=mock_stakick_response), + ): + assert await airos_device.stakick("01:23:45:67:89:aB") + + +@pytest.mark.asyncio +async def test_ap_corners(airos_device, base_url, mode="ap-ptp"): + """Test device operation.""" + cookie = SimpleCookie() + cookie["session_id"] = "test-cookie" + cookie["AIROS_TOKEN"] = "abc123" + + # --- Prepare fake POST /api/auth response with cookies --- + mock_login_response = MagicMock() + mock_login_response.__aenter__.return_value = mock_login_response + mock_login_response.text = AsyncMock(return_value="{}") + mock_login_response.status = 200 + mock_login_response.cookies = cookie + mock_login_response.headers = {"X-CSRF-ID": "test-csrf-token"} + + with ( + patch.object(airos_device.session, "post", return_value=mock_login_response), + patch.object(airos_device, "_use_json_for_login_post", return_value=True), + ): + assert await airos_device.login() + + mock_login_response.cookies = {} + with ( + patch.object(airos_device.session, "post", return_value=mock_login_response), + ): + try: + assert await airos_device.login() + assert False + except airos.exceptions.ConnectionSetupError: + assert True + + mock_login_response.cookies = cookie + mock_login_response.headers = {} + with ( + patch.object(airos_device.session, "post", return_value=mock_login_response), + ): + result = await airos_device.login() + assert result is None + + mock_login_response.headers = {"X-CSRF-ID": "test-csrf-token"} + mock_login_response.text = AsyncMock(return_value="abc123") + with ( + patch.object(airos_device.session, "post", return_value=mock_login_response), + ): + try: + assert await airos_device.login() + assert False + except airos.exceptions.DataMissingError: + assert True + + mock_login_response.text = AsyncMock(return_value="{}") + mock_login_response.status = 400 + with ( + patch.object(airos_device.session, "post", return_value=mock_login_response), + ): + try: + assert await airos_device.login() + assert False + except airos.exceptions.ConnectionAuthenticationError: + assert True + + mock_login_response.status = 200 + with patch.object(airos_device.session, "post", side_effect=aiohttp.ClientError): + try: + assert await airos_device.login() + assert False + except airos.exceptions.DeviceConnectionError: + assert True From d70288bada6057c4e97b2150c595938261cb283a Mon Sep 17 00:00:00 2001 From: Tom Scholten Date: Sat, 19 Jul 2025 20:14:56 +0200 Subject: [PATCH 2/2] CRAI suggestions --- airos/airos8.py | 4 ++-- tests/test_stations.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/airos/airos8.py b/airos/airos8.py index 345ae10..7eaa82b 100644 --- a/airos/airos8.py +++ b/airos/airos8.py @@ -4,7 +4,7 @@ import json import logging -from urllib.parse import urlparse +from urllib.parse import quote, urlparse import aiohttp @@ -235,7 +235,7 @@ async def stakick(self, mac_address: str = None) -> bool: kick_payload = { "staif": "ath0", - "staid": urlparse(mac_address.upper()), + "staid": quote(mac_address.upper(), safe=""), } kick_request_headers["Content-Type"] = ( diff --git a/tests/test_stations.py b/tests/test_stations.py index 7758074..259ee34 100644 --- a/tests/test_stations.py +++ b/tests/test_stations.py @@ -70,7 +70,7 @@ async def test_reconnect(airos_device, base_url): with ( patch.object(airos_device.session, "post", return_value=mock_stakick_response), - patch.object(airos_device, "connected", return_value=mock_stakick_response), + patch.object(airos_device, "connected", True), ): assert await airos_device.stakick("01:23:45:67:89:aB")