Skip to content

Commit f0495ef

Browse files
CoMPaTechCoMPaTech
committed
Improve session handling and ssl, bump version
1 parent b2ead4f commit f0495ef

File tree

2 files changed

+121
-116
lines changed

2 files changed

+121
-116
lines changed

‎airos/airos8.py‎

Lines changed: 120 additions & 114 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,8 @@
22

33
from __future__ import annotations
44

5-
import asyncio
65
import json
76
import logging
8-
import ssl
97

108
import aiohttp
119

@@ -17,12 +15,22 @@
1715
class AirOS8:
1816
"""Set up connection to AirOS."""
1917

20-
def __init__(self, host: str, username: str, password: str):
18+
def __init__(
19+
self,
20+
host: str,
21+
username: str,
22+
password: str,
23+
session: aiohttp.ClientSession,
24+
verify_ssl: bool = True,
25+
):
2126
"""Initialize AirOS8 class."""
2227
self.username = username
2328
self.password = password
2429
self.base_url = f"https://{host}"
2530

31+
self.session = session
32+
self.verify_ssl = verify_ssl
33+
2634
self._login_url = f"{self.base_url}/api/auth" # AirOS 8
2735
self._status_cgi_url = f"{self.base_url}/status.cgi" # AirOS 8
2836
self.current_csrf_token = None
@@ -45,120 +53,116 @@ def __init__(self, host: str, username: str, password: str):
4553

4654
async def login(self) -> bool:
4755
"""Log in to the device assuring cookies and tokens set correctly."""
48-
loop = asyncio.get_running_loop()
49-
ssl_context = await loop.run_in_executor(None, ssl.create_default_context)
50-
ssl_context.check_hostname = False
51-
ssl_context.verify_mode = ssl.CERT_NONE
52-
connector = aiohttp.TCPConnector(ssl=ssl_context)
53-
54-
async with aiohttp.ClientSession(connector=connector) as self.session:
55-
# --- Step 0: Pre-inject the 'ok=1' cookie before login POST (mimics curl) ---
56-
self.session.cookie_jar.update_cookies({"ok": "1"})
57-
58-
# --- Step 1: Attempt Login to /api/auth (This now sets all session cookies and the CSRF token) ---
59-
login_payload = {
60-
"username": self.username,
61-
"password": self.password,
62-
}
63-
64-
login_request_headers = {**self._common_headers}
65-
66-
post_data = None
67-
if self._use_json_for_login_post:
68-
login_request_headers["Content-Type"] = "application/json"
69-
post_data = json.dumps(login_payload)
70-
else:
71-
login_request_headers["Content-Type"] = (
72-
"application/x-www-form-urlencoded; charset=UTF-8"
73-
)
74-
post_data = login_payload
75-
76-
try:
77-
async with self.session.post(
78-
self._login_url, data=post_data, headers=login_request_headers
79-
) as response:
80-
if not response.cookies:
81-
logger.exception("Empty cookies after login, bailing out.")
82-
raise DataMissingError from None
83-
else:
84-
for _, morsel in response.cookies.items():
85-
# If the AIROS_ cookie was parsed but isn't automatically added to the jar, add it manually
86-
if (
87-
morsel.key.startswith("AIROS_")
88-
and morsel.key not in self.session.cookie_jar
89-
):
90-
# `SimpleCookie`'s Morsel objects are designed to be compatible with cookie jars.
91-
# We need to set the domain if it's missing, otherwise the cookie might not be sent.
92-
# For IP addresses, the domain is typically blank.
93-
# aiohttp's jar should handle it, but for explicit control:
94-
if not morsel.get("domain"):
95-
morsel["domain"] = (
96-
response.url.host
97-
) # Set to the host that issued it
98-
self.session.cookie_jar.update_cookies(
99-
{
100-
morsel.key: morsel.output(header="")[
101-
len(morsel.key) + 1 :
102-
]
103-
.split(";")[0]
104-
.strip()
105-
},
106-
response.url,
107-
)
108-
# The update_cookies method can take a SimpleCookie morsel directly or a dict.
109-
# The morsel.output method gives 'NAME=VALUE; Path=...; HttpOnly'
110-
# We just need 'NAME=VALUE' or the morsel object itself.
111-
# Let's use the morsel directly which is more robust.
112-
# Alternatively: self.session.cookie_jar.update_cookies({morsel.key: morsel.value}) might work if it's simpler.
113-
# Aiohttp's update_cookies takes a dict mapping name to value.
114-
# To pass the full morsel with its attributes, we need to add it to the jar's internal structure.
115-
# Simpler: just ensure the key-value pair is there for simple jar.
116-
117-
# Let's try the direct update of the key-value
118-
self.session.cookie_jar.update_cookies(
119-
{morsel.key: morsel.value}
120-
)
121-
122-
new_csrf_token = response.headers.get("X-CSRF-ID")
123-
if new_csrf_token:
124-
self.current_csrf_token = new_csrf_token
125-
else:
126-
return
127-
128-
# Re-check cookies in self.session.cookie_jar AFTER potential manual injection
129-
airos_cookie_found = False
130-
ok_cookie_found = False
131-
if not self.session.cookie_jar:
132-
logger.exception(
133-
"COOKIE JAR IS EMPTY after login POST. This is a major issue."
134-
)
135-
raise DataMissingError from None
136-
for cookie in self.session.cookie_jar:
137-
if cookie.key.startswith("AIROS_"):
138-
airos_cookie_found = True
139-
if cookie.key == "ok":
140-
ok_cookie_found = True
56+
# --- Step 0: Pre-inject the 'ok=1' cookie before login POST (mimics curl) ---
57+
self.session.cookie_jar.update_cookies({"ok": "1"})
14158

142-
if not airos_cookie_found and not ok_cookie_found:
143-
raise DataMissingError from None
59+
# --- Step 1: Attempt Login to /api/auth (This now sets all session cookies and the CSRF token) ---
60+
login_payload = {
61+
"username": self.username,
62+
"password": self.password,
63+
}
14464

145-
response_text = await response.text()
65+
login_request_headers = {**self._common_headers}
14666

147-
if response.status == 200:
148-
try:
149-
json.loads(response_text)
150-
return True
151-
except json.JSONDecodeError as err:
152-
logger.exception("JSON Decode Error")
153-
raise DataMissingError from err
67+
post_data = None
68+
if self._use_json_for_login_post:
69+
login_request_headers["Content-Type"] = "application/json"
70+
post_data = json.dumps(login_payload)
71+
else:
72+
login_request_headers["Content-Type"] = (
73+
"application/x-www-form-urlencoded; charset=UTF-8"
74+
)
75+
post_data = login_payload
15476

155-
else:
156-
log = f"Login failed with status {response.status}. Full Response: {response.text}"
157-
logger.error(log)
158-
raise ConnectionFailedError from None
159-
except aiohttp.ClientError as err:
160-
logger.exception("Error during login")
161-
raise ConnectionFailedError from err
77+
try:
78+
async with self.session.post(
79+
self._login_url,
80+
data=post_data,
81+
headers=login_request_headers,
82+
ssl=self.verify_ssl,
83+
) as response:
84+
if not response.cookies:
85+
logger.exception("Empty cookies after login, bailing out.")
86+
raise DataMissingError from None
87+
else:
88+
for _, morsel in response.cookies.items():
89+
# If the AIROS_ cookie was parsed but isn't automatically added to the jar, add it manually
90+
if (
91+
morsel.key.startswith("AIROS_")
92+
and morsel.key not in self.session.cookie_jar
93+
):
94+
# `SimpleCookie`'s Morsel objects are designed to be compatible with cookie jars.
95+
# We need to set the domain if it's missing, otherwise the cookie might not be sent.
96+
# For IP addresses, the domain is typically blank.
97+
# aiohttp's jar should handle it, but for explicit control:
98+
if not morsel.get("domain"):
99+
morsel["domain"] = (
100+
response.url.host
101+
) # Set to the host that issued it
102+
self.session.cookie_jar.update_cookies(
103+
{
104+
morsel.key: morsel.output(header="")[
105+
len(morsel.key) + 1 :
106+
]
107+
.split(";")[0]
108+
.strip()
109+
},
110+
response.url,
111+
)
112+
# The update_cookies method can take a SimpleCookie morsel directly or a dict.
113+
# The morsel.output method gives 'NAME=VALUE; Path=...; HttpOnly'
114+
# We just need 'NAME=VALUE' or the morsel object itself.
115+
# Let's use the morsel directly which is more robust.
116+
# Alternatively: self.session.cookie_jar.update_cookies({morsel.key: morsel.value}) might work if it's simpler.
117+
# Aiohttp's update_cookies takes a dict mapping name to value.
118+
# To pass the full morsel with its attributes, we need to add it to the jar's internal structure.
119+
# Simpler: just ensure the key-value pair is there for simple jar.
120+
121+
# Let's try the direct update of the key-value
122+
self.session.cookie_jar.update_cookies(
123+
{morsel.key: morsel.value}
124+
)
125+
126+
new_csrf_token = response.headers.get("X-CSRF-ID")
127+
if new_csrf_token:
128+
self.current_csrf_token = new_csrf_token
129+
else:
130+
return
131+
132+
# Re-check cookies in self.session.cookie_jar AFTER potential manual injection
133+
airos_cookie_found = False
134+
ok_cookie_found = False
135+
if not self.session.cookie_jar:
136+
logger.exception(
137+
"COOKIE JAR IS EMPTY after login POST. This is a major issue."
138+
)
139+
raise DataMissingError from None
140+
for cookie in self.session.cookie_jar:
141+
if cookie.key.startswith("AIROS_"):
142+
airos_cookie_found = True
143+
if cookie.key == "ok":
144+
ok_cookie_found = True
145+
146+
if not airos_cookie_found and not ok_cookie_found:
147+
raise DataMissingError from None
148+
149+
response_text = await response.text()
150+
151+
if response.status == 200:
152+
try:
153+
json.loads(response_text)
154+
return True
155+
except json.JSONDecodeError as err:
156+
logger.exception("JSON Decode Error")
157+
raise DataMissingError from err
158+
159+
else:
160+
log = f"Login failed with status {response.status}. Full Response: {response.text}"
161+
logger.error(log)
162+
raise ConnectionFailedError from None
163+
except aiohttp.ClientError as err:
164+
logger.exception("Error during login")
165+
raise ConnectionFailedError from err
162166

163167
async def status(self) -> dict:
164168
"""Retrieve status from the device."""
@@ -169,7 +173,9 @@ async def status(self) -> dict:
169173

170174
try:
171175
async with self.session.get(
172-
self._status_cgi_url, headers=authenticated_get_headers
176+
self._status_cgi_url,
177+
headers=authenticated_get_headers,
178+
ssl=self.verify_ssl,
173179
) as response:
174180
status_response_text = await response.text()
175181

‎pyproject.toml‎

Lines changed: 1 addition & 2 deletions
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.0.3a2"
7+
version = "0.0.4a0"
88
license = "MIT"
99
description = "Ubiquity airOS module(s) for Python 3."
1010
readme = "README.md"
@@ -22,7 +22,6 @@ maintainers = [
2222
requires-python = ">=3.13"
2323
dependencies = [
2424
"aiohttp",
25-
"munch",
2625
]
2726

2827
[project.urls]

0 commit comments

Comments
 (0)