ForcePilot/backend/test/unit/channels/test_wechat_crypto.py

from __future__ import annotations

import hashlib

import pytest

from yuxi.channels.adapters.wechat.mp.crypto import (
    decrypt_message,
    verify_signature,
    verify_url_echostr,
)
from yuxi.channels.adapters.wechat.wecom.crypto import (
    decrypt_message as wecom_decrypt_message,
)
from yuxi.channels.adapters.wechat.wecom.crypto import (
    encrypt_message as wecom_encrypt_message,
)
from yuxi.channels.adapters.wechat.wecom.crypto import (
    verify_url_signature as wecom_verify_url_signature,
)


class TestMpCrypto:
    def test_verify_signature_valid(self):
        token = "test_token"
        timestamp = "1234567890"
        nonce = "random_nonce"
        params = sorted([token, timestamp, nonce])
        expected = hashlib.sha1("".join(params).encode()).hexdigest()
        assert verify_signature(token, timestamp, nonce, expected) is True

    def test_verify_signature_invalid(self):
        assert verify_signature("token", "ts", "nonce", "bad_signature") is False

    def test_verify_signature_wrong_token(self):
        token = "right_token"
        timestamp = "1234567890"
        nonce = "random_nonce"
        params = sorted([token, timestamp, nonce])
        correct_sig = hashlib.sha1("".join(params).encode()).hexdigest()
        assert verify_signature("wrong_token", timestamp, nonce, correct_sig) is False

    def test_verify_url_echostr_valid(self):
        token = "my_token"
        timestamp = "123"
        nonce = "abc"
        params = sorted([token, timestamp, nonce])
        sig = hashlib.sha1("".join(params).encode()).hexdigest()
        ok, result = verify_url_echostr(token, timestamp, nonce, "echo_test", sig)
        assert ok is True
        assert result == "echo_test"

    def test_verify_url_echostr_invalid_signature(self):
        ok, result = verify_url_echostr("token", "ts", "nonce", "echo", "bad")
        assert ok is False
        assert result == ""

    def test_decrypt_message_mp(self):
        import base64
        import os

        from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

        encoding_aes_key_raw = b"\x00" * 32
        encoding_aes_key = base64.b64encode(encoding_aes_key_raw).decode().rstrip("=")
        key = base64.b64decode(encoding_aes_key + "=")
        app_id = "wx_test_app"

        content = "<xml>hello</xml>"
        content_bytes = content.encode("utf-8")
        import struct
        random_bytes = os.urandom(16)
        msg_len = struct.pack("!I", len(content_bytes))
        raw = random_bytes + msg_len + content_bytes + app_id.encode("utf-8")

        pad_len = 32 - (len(raw) % 32)
        if pad_len == 0:
            pad_len = 32
        padded = raw + bytes([pad_len] * pad_len)

        cipher = Cipher(algorithms.AES(key), modes.CBC(key[:16]))
        encryptor = cipher.encryptor()
        ciphertext = encryptor.update(padded) + encryptor.finalize()
        encrypted_msg = base64.b64encode(ciphertext).decode()

        result = decrypt_message(encrypted_msg, encoding_aes_key, app_id)
        assert result == content


class TestWeComCrypto:
    def test_verify_signature_valid(self):
        from yuxi.channels.adapters.wechat.wecom.crypto import verify_signature as wecom_verify
        token = "wecom_token"
        timestamp = "1234567890"
        nonce = "random"
        params = sorted([token, timestamp, nonce])
        expected = hashlib.sha1("".join(params).encode()).hexdigest()
        assert wecom_verify(token, timestamp, nonce, expected) is True

    def test_encrypt_decrypt_roundtrip(self):
        import base64

        encoding_aes_key_raw = b"\x01" * 32
        encoding_aes_key = base64.b64encode(encoding_aes_key_raw).decode().rstrip("=")
        app_id = "test_corp_id"

        content = "测试消息内容"
        encrypted = wecom_encrypt_message(content, encoding_aes_key, app_id)
        decrypted, receive_id = wecom_decrypt_message(encrypted, encoding_aes_key)
        assert decrypted == content
        assert receive_id == app_id

    def test_decrypt_message(self):
        import base64
        import os
        import struct

        from cryptography.hazmat.primitives import padding as sym_padding
        from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

        encoding_aes_key_raw = b"\x02" * 32
        encoding_aes_key = base64.b64encode(encoding_aes_key_raw).decode().rstrip("=")
        key = base64.b64decode(encoding_aes_key + "=")

        content = "解密测试"
        content_bytes = content.encode("utf-8")
        random_bytes = os.urandom(16)
        msg_len = struct.pack("!I", len(content_bytes))
        app_id_bytes = "corp_id".encode("utf-8")
        raw = random_bytes + msg_len + content_bytes + app_id_bytes

        padder = sym_padding.PKCS7(128).padder()
        padded = padder.update(raw) + padder.finalize()

        cipher = Cipher(algorithms.AES(key), modes.CBC(key[:16]))
        encryptor = cipher.encryptor()
        ciphertext = encryptor.update(padded) + encryptor.finalize()
        encrypted = base64.b64encode(ciphertext).decode()

        result, rid = wecom_decrypt_message(encrypted, encoding_aes_key)
        assert result == content
        assert rid == "corp_id"

    def test_verify_url_signature_valid(self):
        token = "url_token"
        timestamp = "111"
        nonce = "222"
        params = sorted([token, timestamp, nonce])
        sig = hashlib.sha1("".join(params).encode()).hexdigest()

        import base64
        echostr = base64.b64encode(b"ok_response").decode()
        ok, result = wecom_verify_url_signature(token, timestamp, nonce, echostr, sig)
        assert ok is True
        assert result == "ok_response"

    def test_verify_url_signature_invalid(self):
        ok, result = wecom_verify_url_signature("t", "1", "2", "echo", "bad_sig")
        assert ok is False
        assert result == ""