diff --git a/app/modules/ugreen/api.py b/app/modules/ugreen/api.py index da576198..1e1f12ed 100644 --- a/app/modules/ugreen/api.py +++ b/app/modules/ugreen/api.py @@ -1,7 +1,7 @@ import base64 import uuid from dataclasses import dataclass -from typing import Any, Dict, Mapping, Optional, Union +from typing import Any, Mapping, Optional, Union from urllib.parse import urlsplit, urlunsplit from requests import Session @@ -13,6 +13,10 @@ from app.utils.url import UrlUtils @dataclass class ApiResult: + """ + 绿联接口标准响应封装。 + """ + code: int = -1 msg: str = "" data: Any = None @@ -21,6 +25,7 @@ class ApiResult: @property def success(self) -> bool: + """判断绿联接口是否返回成功状态""" return self.code == 200 @@ -59,7 +64,17 @@ class Api: ug_agent: str = "PC/WEB", timeout: int = 20, verify_ssl: bool = True, - ): + ) -> None: + """ + 初始化绿联影视 API 客户端。 + + :param host: 绿联服务端地址 + :param client_version: 绿联 Web 客户端版本号 + :param language: 请求语言 + :param ug_agent: 绿联客户端标识 + :param timeout: HTTP 请求超时时间 + :param verify_ssl: 是否校验 HTTPS 证书 + """ self._host = self._normalize_base_url(host) self._session = Session() @@ -80,25 +95,30 @@ class Api: @property def host(self) -> str: + """获取规范化后的绿联服务端地址""" return self._host @property def token(self) -> Optional[str]: + """获取当前登录会话 token""" return self._token @property def static_token(self) -> Optional[str]: + """获取可用于静态资源访问的 token""" return self._static_token @property def is_ugk(self) -> bool: + """判断当前会话是否使用 ugk 访问参数""" return self._is_ugk @property def public_key(self) -> Optional[str]: + """获取当前会话加密公钥""" return self._public_key - def close(self): + def close(self) -> None: """ 关闭底层 HTTP 会话。 """ @@ -262,14 +282,32 @@ class Api: logger.error(f"绿联登录失败：{login_result.msg}") return None - token = str(login_result.data.get("token") or "").strip() - public_key = self._decode_public_key(str(login_result.data.get("public_key") or "")) + token = str( + login_result.data.get("token") + or login_result.data.get("token_id") + or login_result.data.get("tokenId") + or "" + ).strip() + public_key = ( + self._decode_public_key( + str( + login_result.data.get("public_key") + or login_result.data.get("publicKey") + or "" + ) + ) + or login_public_key + ) if not token or not public_key: - logger.error("绿联登录失败：未返回 token/public_key") + logger.error("绿联登录失败：未返回 token/token_id 或可用公钥") return None self._token = token - static_token = str(login_result.data.get("static_token") or "").strip() + static_token = str( + login_result.data.get("static_token") + or login_result.data.get("staticToken") + or "" + ).strip() self._static_token = static_token or self._token self._is_ugk = bool(login_result.data.get("is_ugk")) self._public_key = public_key @@ -365,7 +403,7 @@ class Api: ) return True - def logout(self): + def logout(self) -> None: """ 登出并清理本地认证状态。 """ @@ -569,7 +607,7 @@ class Api: """ 获取海报墙文件夹与条目（可按目录路径递归展开）。 """ - params: Dict[str, Any] = { + params: dict[str, Any] = { "page": page, "page_size": page_size, "sort_type": sort_type, @@ -590,7 +628,7 @@ class Api: """ 获取电影详情。 """ - params: Dict[str, Any] = { + params: dict[str, Any] = { "id": item_id, "media_lib_set_id": media_lib_set_id, "fileVersion": "true", diff --git a/tests/test_ugreen_api.py b/tests/test_ugreen_api.py index cd3c5171..54cc17c4 100644 --- a/tests/test_ugreen_api.py +++ b/tests/test_ugreen_api.py @@ -1,21 +1,28 @@ -import unittest from types import SimpleNamespace +from typing import Optional from unittest.mock import patch from app.modules.ugreen.api import Api class _FakeResponse: - def __init__(self, payload: dict, headers: dict | None = None): + def __init__(self, payload: dict, headers: Optional[dict] = None) -> None: + """初始化伪造 HTTP 响应""" self._payload = payload self.headers = headers or {} - def json(self): + def json(self) -> dict: + """返回伪造 JSON 响应体""" return self._payload class _FakeSession: - def __init__(self, get_responses=None, post_responses=None): + def __init__( + self, + get_responses: Optional[list[_FakeResponse]] = None, + post_responses: Optional[list[_FakeResponse]] = None, + ) -> None: + """初始化伪造 HTTP 会话""" self._get_responses = list(get_responses or []) self._post_responses = list(post_responses or []) self.calls: list[tuple[str, dict]] = [] @@ -24,86 +31,149 @@ class _FakeSession: update=lambda *_args, **_kwargs: None, ) - def get(self, *args, **kwargs): + def get(self, *args: object, **kwargs: object) -> _FakeResponse: + """记录 GET 请求并返回预置响应""" if args: kwargs = {"url": args[0], **kwargs} self.calls.append(("GET", kwargs)) return self._get_responses.pop(0) if self._get_responses else _FakeResponse({}) - def post(self, *args, **kwargs): + def post(self, *args: object, **kwargs: object) -> _FakeResponse: + """记录 POST 请求并返回预置响应""" if args: kwargs = {"url": args[0], **kwargs} self.calls.append(("POST", kwargs)) return self._post_responses.pop(0) if self._post_responses else _FakeResponse({}) @staticmethod - def close(): + def close() -> None: + """关闭伪造会话""" return None class _FakeCrypto: - def __init__(self, *args, **kwargs): - pass + def __init__(self, *args: object, **kwargs: object) -> None: + """初始化伪造加密工具""" + return None @staticmethod def rsa_encrypt_long(raw: str) -> str: + """返回可断言的伪造 RSA 加密内容""" return f"enc:{raw}" @staticmethod - def build_encrypted_request(url: str, method: str = "GET", params=None, **kwargs): - return SimpleNamespace(url=url, headers={}, params=params or {}, json=None, aes_key="k") + def build_encrypted_request( + url: str, + method: str = "GET", + params: Optional[dict] = None, + **kwargs: object, + ) -> SimpleNamespace: + """构造伪造加密请求对象""" + _ = method, kwargs + return SimpleNamespace( + url=url, + headers={}, + params=params or {}, + json=None, + aes_key="k", + ) @staticmethod - def decrypt_response(payload, aes_key): + def decrypt_response(payload: dict, aes_key: str) -> dict: + """原样返回伪造响应内容""" + _ = aes_key return payload -class UgreenApiVerifySslTest(unittest.TestCase): - def test_request_json_default_verify_ssl_true(self): - api = Api(host="https://example.com") - fake_session = _FakeSession( - get_responses=[_FakeResponse({"code": 200})], - post_responses=[_FakeResponse({"code": 200})], - ) - api._session = fake_session +def test_request_json_default_verify_ssl_true() -> None: + """默认请求应开启 HTTPS 证书校验""" + api = Api(host="https://example.com") + fake_session = _FakeSession( + get_responses=[_FakeResponse({"code": 200})], + post_responses=[_FakeResponse({"code": 200})], + ) + api._session = fake_session - api._request_json(url="https://example.com/a", method="GET") - api._request_json(url="https://example.com/b", method="POST", json_data={"x": 1}) + api._request_json(url="https://example.com/a", method="GET") + api._request_json(url="https://example.com/b", method="POST", json_data={"x": 1}) - self.assertEqual(fake_session.calls[0][1].get("verify"), True) - self.assertEqual(fake_session.calls[1][1].get("verify"), True) + assert fake_session.calls[0][1].get("verify") is True + assert fake_session.calls[1][1].get("verify") is True - def test_login_logout_follow_verify_ssl_flag(self): - api = Api(host="https://example.com", verify_ssl=False) - fake_session = _FakeSession( - get_responses=[_FakeResponse({})], - post_responses=[ - _FakeResponse({"code": 200, "msg": "ok", "data": {}}, headers={"x-rsa-token": "BEGIN TEST"}), - _FakeResponse( - { - "code": 200, - "msg": "ok", - "data": { - "token": "token-value", - "public_key": "BEGIN LOGIN KEY", - "static_token": "static-token", - "is_ugk": False, - }, - } - ), - ], - ) - api._session = fake_session - with patch("app.modules.ugreen.api.UgreenCrypto", _FakeCrypto): - token = api.login("tester", "pwd") - self.assertEqual(token, "token-value") - api.logout() +def test_login_logout_follow_verify_ssl_flag() -> None: + """登录与登出请求应沿用用户配置的证书校验开关""" + api = Api(host="https://example.com", verify_ssl=False) + fake_session = _FakeSession( + get_responses=[_FakeResponse({})], + post_responses=[ + _FakeResponse( + {"code": 200, "msg": "ok", "data": {}}, + headers={"x-rsa-token": "BEGIN CHECK KEY"}, + ), + _FakeResponse( + { + "code": 200, + "msg": "ok", + "data": { + "token": "token-value", + "public_key": "BEGIN LOGIN KEY", + "static_token": "static-token", + "is_ugk": False, + }, + } + ), + ], + ) + api._session = fake_session - self.assertEqual(len(fake_session.calls), 3) - self.assertEqual(fake_session.calls[0][0], "POST") - self.assertEqual(fake_session.calls[1][0], "POST") - self.assertEqual(fake_session.calls[2][0], "GET") - self.assertEqual(fake_session.calls[0][1].get("verify"), False) - self.assertEqual(fake_session.calls[1][1].get("verify"), False) - self.assertEqual(fake_session.calls[2][1].get("verify"), False) + with patch("app.modules.ugreen.api.UgreenCrypto", _FakeCrypto): + token = api.login("tester", "pwd") + assert token == "token-value" + assert api.public_key == "BEGIN LOGIN KEY" + assert api.static_token == "static-token" + api.logout() + + assert len(fake_session.calls) == 3 + assert fake_session.calls[0][0] == "POST" + assert fake_session.calls[1][0] == "POST" + assert fake_session.calls[2][0] == "GET" + assert fake_session.calls[0][1].get("verify") is False + assert fake_session.calls[1][1].get("verify") is False + assert fake_session.calls[2][1].get("verify") is False + + +def test_login_accepts_token_id_and_reuses_check_public_key() -> None: + """登录响应只有 token_id 时应复用检查接口公钥完成会话初始化""" + api = Api(host="https://example.com") + fake_session = _FakeSession( + post_responses=[ + _FakeResponse( + {"code": 200, "msg": "ok", "data": {}}, + headers={"x-rsa-token": "BEGIN CHECK KEY"}, + ), + _FakeResponse( + { + "code": 200, + "msg": "success", + "data": { + "enable_otp": True, + "is_exceed": False, + "role": "admin", + "token_id": "token-id-value", + "uid": 1000, + "urgent_email": "tester@example.com", + }, + } + ), + ], + ) + api._session = fake_session + + with patch("app.modules.ugreen.api.UgreenCrypto", _FakeCrypto): + token = api.login("tester", "pwd") + + assert token == "token-id-value" + assert api.token == "token-id-value" + assert api.static_token == "token-id-value" + assert api.public_key == "BEGIN CHECK KEY"