fix: support ugreen token_id login response

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",

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"