Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -10,3 +10,4 @@ tests/__pycache__
.vscode
.coverage
tmp
todo
59 changes: 53 additions & 6 deletions airos/airos8.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

import json
import logging
from urllib.parse import urlparse
from urllib.parse import quote, urlparse

import aiohttp

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()

Expand All @@ -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

Expand Down Expand Up @@ -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": quote(mac_address.upper(), safe=""),
}

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
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
87 changes: 87 additions & 0 deletions tests/test_stations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Expand Down Expand Up @@ -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", True),
):
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