99 lines
4.0 KiB
Python
99 lines
4.0 KiB
Python
"""core/wechat/wecom_crypto.py 单测 —— 企微回调加解密(WXBizMsgCrypt 等价实现)。
|
|
|
|
模块只做入站解密+验签(出站走主动推不需要加密),测试侧按同一方案**自造密文**:
|
|
key = b64decode(EncodingAESKey+'='),AES-256-CBC(IV=key[:16]),明文体 =
|
|
random(16) || len(4B 大端) || msg || receiveid,PKCS7 pad 到 32。
|
|
凭据走 mock.patch.dict 注入合成值,不碰真实 env / 网络。
|
|
此前 wechat 子系统零测试(架构审查 Top3),crypto 是其中唯一可纯测的一块。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import os
|
|
import secrets
|
|
import struct
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
|
|
|
from core.wechat import wecom_crypto as wc
|
|
|
|
_TOKEN = "testtoken"
|
|
# 43 字符合法 EncodingAESKey(b64 去掉尾部 '=')
|
|
_AESKEY = base64.b64encode(secrets.token_bytes(32)).decode().rstrip("=")
|
|
_CORPID = "ww1234567890abcdef"
|
|
_ENV = {"WECOM_CALLBACK_TOKEN": _TOKEN, "WECOM_CALLBACK_AESKEY": _AESKEY}
|
|
|
|
|
|
def _encrypt(msg: str, receiveid: str = _CORPID) -> str:
|
|
"""按企业微信方案造密文(测试侧加密,与被测解密互逆)。"""
|
|
key = base64.b64decode(_AESKEY + "=")
|
|
body = secrets.token_bytes(16) + struct.pack(">I", len(msg.encode())) \
|
|
+ msg.encode() + receiveid.encode()
|
|
pad = 32 - (len(body) % 32) or 32
|
|
body += bytes([pad]) * pad
|
|
enc = Cipher(algorithms.AES(key), modes.CBC(key[:16])).encryptor()
|
|
return base64.b64encode(enc.update(body) + enc.finalize()).decode()
|
|
|
|
|
|
class WecomCryptoTests(unittest.TestCase):
|
|
def setUp(self):
|
|
patcher = mock.patch.dict(os.environ, _ENV)
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
|
|
def test_configured_gate(self):
|
|
self.assertTrue(wc.callback_configured())
|
|
with mock.patch.dict(os.environ, {"WECOM_CALLBACK_TOKEN": ""}):
|
|
self.assertFalse(wc.callback_configured())
|
|
|
|
def test_verify_url_roundtrip(self):
|
|
echostr_plain = "echo-123-测试"
|
|
enc = _encrypt(echostr_plain)
|
|
sig = wc._signature("111", "222", enc)
|
|
self.assertEqual(wc.verify_url(sig, "111", "222", enc, corpid=_CORPID), echostr_plain)
|
|
|
|
def test_verify_url_bad_signature(self):
|
|
enc = _encrypt("x")
|
|
with self.assertRaises(ValueError):
|
|
wc.verify_url("deadbeef", "111", "222", enc, corpid=_CORPID)
|
|
|
|
def test_decrypt_message_roundtrip(self):
|
|
plain_xml = (
|
|
"<xml><FromUserName>zhangsan</FromUserName><MsgType>text</MsgType>"
|
|
"<Content>你好 世界</Content><MsgId>42</MsgId></xml>"
|
|
)
|
|
enc = _encrypt(plain_xml)
|
|
sig = wc._signature("111", "222", enc)
|
|
envelope = f"<xml><Encrypt>{enc}</Encrypt></xml>"
|
|
msg = wc.decrypt_message(sig, "111", "222", envelope, corpid=_CORPID)
|
|
self.assertEqual(msg["FromUserName"], "zhangsan")
|
|
self.assertEqual(msg["MsgType"], "text")
|
|
self.assertEqual(msg["Content"], "你好 世界")
|
|
|
|
def test_receiveid_mismatch_rejected(self):
|
|
enc = _encrypt("<xml><MsgType>text</MsgType></xml>", receiveid="ww_other_corp")
|
|
sig = wc._signature("111", "222", enc)
|
|
with self.assertRaises(ValueError):
|
|
wc.decrypt_message(sig, "111", "222", f"<xml><Encrypt>{enc}</Encrypt></xml>",
|
|
corpid=_CORPID)
|
|
|
|
def test_bad_padding_rejected(self):
|
|
# 用错 key 加密 → 被测侧解出的 padding 大概率非法(1..32 之外)或后续解析崩
|
|
key = secrets.token_bytes(32)
|
|
body = secrets.token_bytes(64)
|
|
enc_obj = Cipher(algorithms.AES(key), modes.CBC(key[:16])).encryptor()
|
|
garbage = base64.b64encode(enc_obj.update(body) + enc_obj.finalize()).decode()
|
|
sig = wc._signature("111", "222", garbage)
|
|
with self.assertRaises(Exception):
|
|
wc.verify_url(sig, "111", "222", garbage, corpid=_CORPID)
|
|
|
|
def test_parse_message_flat_tags(self):
|
|
d = wc.parse_message("<xml><A>1</A><B></B></xml>")
|
|
self.assertEqual(d, {"A": "1", "B": ""})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|