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

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()