新增绿联接口加解密工具与单元测试

app/utils/ugreen_crypto.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import base64
+import hashlib
+import json
+import os
+import uuid
+from dataclasses import dataclass
+from typing import Any, Mapping, Sequence
+from urllib.parse import quote, urlencode, urlsplit, urlunsplit
+
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import padding
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+
+@dataclass
+class UgreenEncryptedRequest:
+    url: str
+    headers: dict[str, str]
+    params: dict[str, str]
+    json: dict[str, Any] | None
+    aes_key: str
+    plain_query: str
+
+
+class UgreenCrypto:
+    """
+    绿联接口请求加解密工具。
+    """
+
+    def __init__(
+        self,
+        public_key: str,
+        token: str | None = None,
+        client_id: str | None = None,
+        client_version: str | None = "76363",
+        ug_agent: str | None = "PC/WEB",
+        language: str = "zh-CN",
+    ) -> None:
+        self.public_key_pem = self.normalize_public_key(public_key)
+        self.public_key = serialization.load_pem_public_key(
+            self.public_key_pem.encode("utf-8")
+        )
+        self.token = token
+        self.client_id = client_id
+        self.client_version = client_version
+        self.ug_agent = ug_agent
+        self.language = language
+
+    @staticmethod
+    def normalize_public_key(public_key: str) -> str:
+        key = (public_key or "").strip().strip('"').replace("\\n", "\n")
+        if "BEGIN" in key:
+            return key if key.endswith("\n") else f"{key}\n"
+        return (
+            "-----BEGIN RSA PUBLIC KEY-----\n"
+            f"{key}\n"
+            "-----END RSA PUBLIC KEY-----\n"
+        )
+
+    @staticmethod
+    def generate_aes_key() -> str:
+        return uuid.uuid4().hex
+
+    @staticmethod
+    def _flatten_query(prefix: str, value: Any) -> list[tuple[str, str]]:
+        pairs: list[tuple[str, str]] = []
+        if isinstance(value, Mapping):
+            for key, item in value.items():
+                next_prefix = f"{prefix}[{key}]" if prefix else str(key)
+                pairs.extend(UgreenCrypto._flatten_query(next_prefix, item))
+            return pairs
+        if isinstance(value, Sequence) and not isinstance(
+            value, (str, bytes, bytearray)
+        ):
+            for item in value:
+                next_prefix = f"{prefix}[]"
+                pairs.extend(UgreenCrypto._flatten_query(next_prefix, item))
+            return pairs
+        if isinstance(value, bool):
+            pairs.append((prefix, "true" if value else "false"))
+            return pairs
+        if value is None:
+            pairs.append((prefix, ""))
+            return pairs
+        pairs.append((prefix, str(value)))
+        return pairs
+
+    @classmethod
+    def encode_query(cls, params: Mapping[str, Any] | None) -> str:
+        if not params:
+            return ""
+        pairs: list[tuple[str, str]] = []
+        for key, value in params.items():
+            pairs.extend(cls._flatten_query(str(key), value))
+        return urlencode(pairs, doseq=False, quote_via=quote, safe="")
+
+    def rsa_encrypt_long(self, plaintext: str) -> str:
+        if not plaintext:
+            return ""
+        key_size = self.public_key.key_size // 8
+        max_chunk = key_size - 11
+        encrypted_chunks: list[bytes] = []
+        raw = plaintext.encode("utf-8")
+        for start in range(0, len(raw), max_chunk):
+            chunk = raw[start : start + max_chunk]
+            encrypted_chunks.append(
+                self.public_key.encrypt(chunk, padding.PKCS1v15())
+            )
+        return base64.b64encode(b"".join(encrypted_chunks)).decode("utf-8")
+
+    @staticmethod
+    def aes_gcm_encrypt(plaintext: str, aes_key: str) -> str:
+        iv = os.urandom(12)
+        cipher = AESGCM(aes_key.encode("utf-8"))
+        encrypted = cipher.encrypt(iv, plaintext.encode("utf-8"), None)
+        # encrypt 返回 ciphertext + tag
+        return base64.b64encode(iv + encrypted).decode("utf-8")
+
+    @staticmethod
+    def aes_gcm_decrypt(payload_b64: str, aes_key: str) -> str:
+        raw = base64.b64decode(payload_b64)
+        iv = raw[:12]
+        encrypted = raw[12:]
+        cipher = AESGCM(aes_key.encode("utf-8"))
+        plain = cipher.decrypt(iv, encrypted, None)
+        return plain.decode("utf-8")
+
+    @staticmethod
+    def build_security_key(token: str) -> str:
+        return hashlib.md5(token.encode("utf-8")).hexdigest()
+
+    @staticmethod
+    def _normalize_body(data: Any) -> str:
+        if isinstance(data, str):
+            return data
+        if isinstance(data, (bytes, bytearray)):
+            return bytes(data).decode("utf-8")
+        return json.dumps(data, ensure_ascii=False, separators=(",", ":"))
+
+    def encrypt_body(self, data: Any, aes_key: str) -> dict[str, str]:
+        plain = self._normalize_body(data)
+        return {
+            "encrypt_req_body": self.aes_gcm_encrypt(plain, aes_key),
+            "req_body_sha256": hashlib.sha256(plain.encode("utf-8")).hexdigest(),
+        }
+
+    def build_headers(
+        self,
+        aes_key: str,
+        token: str | None = None,
+        extra_headers: Mapping[str, str] | None = None,
+        encrypt_token: bool = True,
+    ) -> dict[str, str]:
+        token_value = token if token is not None else self.token
+        headers: dict[str, str] = dict(extra_headers or {})
+
+        if self.client_id:
+            headers.setdefault("Client-Id", self.client_id)
+        if self.client_version:
+            headers.setdefault("Client-Version", self.client_version)
+        if self.ug_agent:
+            headers.setdefault("UG-Agent", self.ug_agent)
+        headers.setdefault("X-Specify-Language", self.language)
+        headers.setdefault("Accept", "application/json, text/plain, */*")
+
+        if token_value:
+            headers["X-Ugreen-Security-Key"] = self.build_security_key(token_value)
+            headers["X-Ugreen-Security-Code"] = self.rsa_encrypt_long(aes_key)
+            headers["X-Ugreen-Token"] = (
+                self.rsa_encrypt_long(token_value) if encrypt_token else token_value
+            )
+        return headers
+
+    def build_encrypted_request(
+        self,
+        url: str,
+        method: str = "GET",
+        params: Mapping[str, Any] | None = None,
+        data: Any | None = None,
+        extra_headers: Mapping[str, str] | None = None,
+        token: str | None = None,
+        encrypt_token: bool = True,
+        encrypt_body: bool = True,
+    ) -> UgreenEncryptedRequest:
+        parsed = urlsplit(url)
+        clean_url = urlunsplit(
+            (parsed.scheme, parsed.netloc, parsed.path, "", parsed.fragment)
+        )
+
+        url_query_plain = parsed.query
+        input_query_plain = self.encode_query(params)
+        plain_query = "&".join(filter(None, [url_query_plain, input_query_plain]))
+
+        aes_key = self.generate_aes_key()
+        encrypted_query = self.aes_gcm_encrypt(plain_query, aes_key)
+
+        req_json = None
+        if data is not None:
+            req_json = self.encrypt_body(data, aes_key) if encrypt_body else data
+
+        headers = self.build_headers(
+            aes_key=aes_key,
+            token=token,
+            extra_headers=extra_headers,
+            encrypt_token=encrypt_token,
+        )
+        if req_json is not None:
+            headers.setdefault("Content-Type", "application/json")
+
+        _ = method  # 保留参数，便于上层统一调用
+
+        return UgreenEncryptedRequest(
+            url=clean_url,
+            headers=headers,
+            params={"encrypt_query": encrypted_query},
+            json=req_json,
+            aes_key=aes_key,
+            plain_query=plain_query,
+        )
+
+    def decrypt_response(self, response_json: Any, aes_key: str) -> Any:
+        if not isinstance(response_json, Mapping):
+            return response_json
+        encrypted = response_json.get("encrypt_resp_body")
+        if not encrypted:
+            return response_json
+        plain = self.aes_gcm_decrypt(str(encrypted), aes_key)
+        try:
+            return json.loads(plain)
+        except json.JSONDecodeError:
+            return plain
+

tests/test_ugreen_crypto.py

@@ -0,0 +1,95 @@
+import base64
+import hashlib
+import json
+import unittest
+
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import padding, rsa
+
+from app.utils.ugreen_crypto import UgreenCrypto
+
+
+def _generate_rsa_keys() -> tuple[str, rsa.RSAPrivateKey]:
+    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    public_pem = private_key.public_key().public_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PublicFormat.PKCS1,
+    ).decode("utf-8")
+    return public_pem, private_key
+
+
+def _rsa_decrypt_long(private_key: rsa.RSAPrivateKey, payload_b64: str) -> str:
+    encrypted = base64.b64decode(payload_b64)
+    chunk_size = private_key.key_size // 8
+    plain_chunks = []
+    for start in range(0, len(encrypted), chunk_size):
+        chunk = encrypted[start : start + chunk_size]
+        plain_chunks.append(private_key.decrypt(chunk, padding.PKCS1v15()))
+    return b"".join(plain_chunks).decode("utf-8")
+
+
+class UgreenCryptoTest(unittest.TestCase):
+    def setUp(self):
+        self.public_key, self.private_key = _generate_rsa_keys()
+        self.token = "demo-token-for-test"
+        self.crypto = UgreenCrypto(
+            public_key=self.public_key,
+            token=self.token,
+            client_id="test-client-id",
+        )
+
+    def test_rsa_encrypt_long(self):
+        plain = "A" * 400
+        encrypted = self.crypto.rsa_encrypt_long(plain)
+        self.assertEqual(plain, _rsa_decrypt_long(self.private_key, encrypted))
+
+    def test_build_encrypted_request_and_decrypt_response(self):
+        req = self.crypto.build_encrypted_request(
+            url="http://127.0.0.1:9999/ugreen/v1/video/homepage/media_list",
+            params={"page": 1, "page_size": 50},
+            data={"foo": "bar", "count": 2},
+        )
+
+        self.assertEqual(
+            req.plain_query,
+            "page=1&page_size=50",
+        )
+        self.assertEqual(
+            req.plain_query,
+            self.crypto.aes_gcm_decrypt(req.params["encrypt_query"], req.aes_key),
+        )
+
+        self.assertEqual(
+            req.headers["X-Ugreen-Security-Key"],
+            hashlib.md5(self.token.encode("utf-8")).hexdigest(),
+        )
+        self.assertEqual(
+            req.aes_key,
+            _rsa_decrypt_long(self.private_key, req.headers["X-Ugreen-Security-Code"]),
+        )
+        self.assertEqual(
+            self.token,
+            _rsa_decrypt_long(self.private_key, req.headers["X-Ugreen-Token"]),
+        )
+
+        encrypted_body = req.json["encrypt_req_body"]
+        body_plain = self.crypto.aes_gcm_decrypt(encrypted_body, req.aes_key)
+        self.assertEqual(json.loads(body_plain), {"foo": "bar", "count": 2})
+        self.assertEqual(
+            req.json["req_body_sha256"],
+            hashlib.sha256(body_plain.encode("utf-8")).hexdigest(),
+        )
+
+        server_payload = {"code": 0, "msg": "ok", "data": {"items": [1, 2, 3]}}
+        resp = {
+            "encrypt_resp_body": self.crypto.aes_gcm_encrypt(
+                json.dumps(server_payload, ensure_ascii=False, separators=(",", ":")),
+                req.aes_key,
+            )
+        }
+        decoded = self.crypto.decrypt_response(resp, req.aes_key)
+        self.assertEqual(decoded, server_payload)
+
+
+if __name__ == "__main__":
+    unittest.main()