chuanqi/tests/test_msg_codec.py

# tests/test_msg_codec.py
"""
cpp_reversed/msg_codec 算法库的 pytest 测试套件
覆盖: hdr_codec / frame / xxtea / msg_codec 顶层 API / 抓包分析场景
"""
import os
import sys
import struct

import pytest

# 让 pytest 能 import cpp_reversed 模块
CPP_REV = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "cpp_reversed"))
if CPP_REV not in sys.path:
    sys.path.insert(0, CPP_REV)

# 兼容 protocol_bot (用于 cross-verify)
PROTO_BOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "protocol_bot"))
if PROTO_BOT not in sys.path:
    sys.path.insert(0, PROTO_BOT)

import hdr_codec
import xxtea
import frame as frame_mod
import msg_codec


# ============== hdr_codec: 24B 头编解码 ==============

class TestHdrCodec:
    """测试 24B 头编解码 (Encode6BitBuf 算法核心)"""

    def test_encode_24b_to_32b(self):
        plain = b"\x00" * 24
        enc = hdr_codec.encode_header(plain)
        assert len(enc) == 32
        assert all(0x3C <= b <= 0x7B for b in enc), "编码字节非可打印"

    def test_decode_32b_to_24b(self):
        plain = b"\x00" * 24
        enc = hdr_codec.encode_header(plain)
        dec = hdr_codec.decode_header(enc)
        assert dec == plain

    def test_roundtrip_random(self):
        import random
        random.seed(42)
        for _ in range(50):
            plain = bytes(random.randint(0, 255) for _ in range(24))
            enc = hdr_codec.encode_header(plain)
            dec = hdr_codec.decode_header(enc)
            assert dec == plain, f"roundtrip fail: {plain.hex()} → {enc.hex()} → {dec.hex()}"

    def test_parse_header_legacy_24b(self):
        # 24B 旧格式: recog(4) + p1(4) + p2(4) + p3(4) + p4(4) + msgId(2) + 2 字节
        hdr = bytearray(24)
        struct.pack_into("<I", hdr, 0, 900000)    # recog
        struct.pack_into("<I", hdr, 4, 12345)     # p1
        struct.pack_into("<I", hdr, 8, 37504)     # p2
        struct.pack_into("<H", hdr, 20, 101)      # msgId
        fields = hdr_codec.parse_header(bytes(hdr))
        assert fields["msgId"] == 101
        assert fields["recog"] == 900000
        assert fields["param1"] == 12345
        assert fields["param2"] == 37504


# ============== frame: 36B 头 + body 完整帧 ==============

class TestFrame:
    """测试 36B wire format 头 + body 完整帧"""

    def test_pack_frame_basic(self):
        packed = frame_mod.pack_frame(101, 900000, 0, 0, 0, 0, b"")
        assert len(packed) == 48  # 48B 头 + 0B body

    def test_pack_with_body(self):
        body = b"hello world" * 10  # 110B
        packed = frame_mod.pack_frame(101, 900000, 0, 0, 0, 0, body)
        assert len(packed) == 48 + len(body)

    def test_roundtrip_basic(self):
        f = frame_mod.Frame(msg_id=101, recog=900000, p1=12345, p2=37504, p3=0, tail=0, body=b"test")
        packed = f.pack()
        frames, rest = frame_mod.unpack_stream(packed)
        assert rest == b""
        assert len(frames) == 1
        g = frames[0]
        assert g.msg_id == 101
        assert g.recog == 900000
        assert g.body == b"test"

    def test_roundtrip_random_50(self):
        import random
        random.seed(0xC0FFEE)
        for i in range(50):
            msg_id = random.randint(0, 0x1FFE)
            recog = random.randint(0, 0x7FFFFFFF)
            p1 = random.randint(0, 0x7FFFFFFF)
            p2 = random.randint(0, 0x7FFFFFFF)
            p3 = random.randint(0, 0x7FFFFFFF)
            tail = random.randint(0, 9)
            body = bytes(random.randint(0, 255) for _ in range(random.randint(0, 200)))
            f = frame_mod.Frame(msg_id, recog, p1, p2, p3, tail, body)
            packed = f.pack()
            frames, rest = frame_mod.unpack_stream(packed)
            assert rest == b"", f"iter {i}: rest non-empty"
            assert len(frames) == 1, f"iter {i}: wrong frame count"
            g = frames[0]
            assert g.msg_id == msg_id
            assert g.recog == recog
            assert g.p1 == p1
            assert g.p2 == p2
            assert g.p3 == p3
            assert g.tail == tail
            assert g.body == body

    def test_unpack_partial(self):
        f = frame_mod.Frame(msg_id=529, recog=0, p1=0, p2=0, p3=0, tail=1, body=b"key32abc")
        packed = f.pack()
        frames, rest = frame_mod.unpack_stream(packed[:-5])
        assert frames == []
        assert rest == packed[:-5]

    def test_unpack_multi_packet(self):
        f1 = frame_mod.Frame(msg_id=20, recog=0, p1=0, p2=0, p3=0, tail=9, body=b"login_payload")
        f2 = frame_mod.Frame(msg_id=529, recog=0, p1=0, p2=0, p3=0, tail=1, body=b"key32abc")
        multi = f1.pack() + f2.pack()
        frames, rest = frame_mod.unpack_stream(multi)
        assert rest == b""
        assert len(frames) == 2
        assert frames[0].msg_id == 20
        assert frames[1].msg_id == 529
        assert frames[1].body == b"key32abc"

    def test_wire_format_interop_with_protocol_bot(self):
        """与 protocol_bot/encode6bitbuf_final 字节级一致 (向后兼容)"""
        from encode6bitbuf_final import build_frame, encode6
        head36 = frame_mod.HDR_STRUCT.pack(0, 101, 900000, 0, 0, 0, 0, 0, 0)
        assert frame_mod._encode_full(head36) == encode6(head36)
        # build_frame 缺省 body=空
        assert frame_mod.pack_frame(101, 900000, 0, 0, 0, 0, b"") == build_frame(101, 900000, 0, 0, 0)


# ============== xxtea: 加解密 ==============

class TestXXTEA:
    """测试 XXTEA 加解密算法"""

    KEY = b"0123456789abcdef"

    def test_known_vector_zero(self):
        # 16B zeros + zero key → e6c8d5ff070fb6e498a534f7ac03e399
        c = xxtea.encrypt(b"\x00" * 16, b"\x00" * 16)
        assert c.hex() == "e6c8d5ff070fb6e498a534f7ac03e399"

    def test_roundtrip_8b(self):
        plain = b"\x12" * 8
        c = xxtea.encrypt(plain, self.KEY)
        d = xxtea.decrypt(c, self.KEY)
        assert d == plain

    def test_roundtrip_16b(self):
        plain = bytes.fromhex("123456789abcdef0deadbeefcafebabe")
        c = xxtea.encrypt(plain, self.KEY)
        d = xxtea.decrypt(c, self.KEY)
        assert d == plain

    def test_roundtrip_32b(self):
        plain = b"Hello XXTEA!!!  "  # 16B
        c = xxtea.encrypt(plain, self.KEY)
        d = xxtea.decrypt(c, self.KEY)
        assert d == plain

    def test_padded_roundtrip(self):
        for plain in [b"", b"x", b"xy", b"Hello", b"HelloWorld12"]:
            enc = xxtea.encrypt_padded(plain, self.KEY)
            dec = xxtea.decrypt_padded(enc, self.KEY)
            assert dec == plain, f"padded roundtrip fail: {plain!r}"

    def test_empty(self):
        c = xxtea.encrypt(b"", self.KEY)
        d = xxtea.decrypt(c, self.KEY)
        assert d == b""

    def test_invalid_key_length(self):
        with pytest.raises(ValueError, match="key 必须是 16 字节"):
            xxtea.encrypt(b"\x00" * 8, b"short")

    def test_invalid_plain_length(self):
        with pytest.raises(ValueError, match="plain 长度必须是 8 倍数"):
            xxtea.encrypt(b"7bytes!", self.KEY)

    def test_cross_verify_with_third_party(self):
        """与第三方 xxtea 6.2.0 包 cross-verify (需先 pip install xxtea)"""
        try:
            import xxtea as third_party
        except ImportError:
            pytest.skip("xxtea 第三方包未安装")
        test_cases = [
            (b"\x00" * 16, b"\x00" * 16),
            (b"\xff" * 16, b"\x00" * 16),
            (bytes.fromhex("123456789abcdef0deadbeefcafebabe"),
             bytes.fromhex("0102030405060708090a0b0c0d0e0f10")),
        ]
        for plain, key in test_cases:
            mine = xxtea.encrypt(plain, key)
            theirs = third_party.encrypt(plain, key)
            assert mine == theirs, f"cross-verify fail: {plain.hex()} / {key.hex()}"


# ============== msg_codec: 顶层 API ==============

class TestMsgCodec:
    """测试 msg_codec 顶层 API"""

    def test_send_basic(self):
        wire = msg_codec.send(101, 900000, 0, 0, 0, b"hello", 0)
        assert len(wire) == 48 + 5

    def test_receive_roundtrip(self):
        wire = msg_codec.send(20, 0, 0, 0, 0, b"*#test/1/2/3", 9)
        frames, rest = msg_codec.receive(wire)
        assert rest == b""
        assert len(frames) == 1
        assert frames[0].msg_id == 20
        assert frames[0].body == b"*#test/1/2/3"

    def test_frame_to_dict(self):
        f = frame_mod.Frame(101, 900000, 12345, 37504, 0, 0, b"hello world!")
        d = msg_codec.frame_to_dict(f)
        assert d["msg_id"] == 101
        assert d["recog"] == 900000
        assert d["body_len"] == 12
        assert d["body_preview"] == b"hello world!".hex()

    def test_decode_captured_48b_wire(self):
        wire = msg_codec.send(101, 900000, 12345, 37504, 0, b"x" * 32)
        d = msg_codec.decode_captured(wire[:48])
        assert d["msgId"] == 101
        assert d["recog"] == 900000
        assert d["p1"] == 12345
        assert d["p2"] == 37504
        assert d["body_len"] == 32

    def test_decode_captured_32b_legacy(self):
        # 构造 24B 旧格式明文头
        plain = bytearray(24)
        struct.pack_into("<I", plain, 0, 900000)
        struct.pack_into("<I", plain, 4, 12345)
        struct.pack_into("<I", plain, 8, 37504)
        struct.pack_into("<H", plain, 20, 101)
        enc = hdr_codec.encode_header(bytes(plain))
        d = msg_codec.decode_captured(enc)
        assert d["msgId"] == 101
        assert d["recog"] == 900000
        assert d["param1"] == 12345
        assert d["param2"] == 37504

    def test_decode_captured_too_short(self):
        with pytest.raises(ValueError, match="至少需要 32B"):
            msg_codec.decode_captured(b"\x00" * 10)

    def test_xxtea_aliases(self):
        c = msg_codec.xxtea_encrypt(b"test1234", self.KEY)
        d = msg_codec.xxtea_decrypt(c, self.KEY)
        assert d == b"test1234"

    def test_padded_aliases(self):
        c = msg_codec.encrypt_padded(b"test", self.KEY)
        d = msg_codec.decrypt_padded(c, self.KEY)
        assert d == b"test"

    @property
    def KEY(self):
        return b"0123456789abcdef"


# ============== rsa: 接口契约 ==============

class TestRSAStub:
    """测试 rsa 接口契约 (密钥未提取, 抛 NotImplementedError)"""

    def test_encrypt_pub_raises(self):
        with pytest.raises(NotImplementedError, match="未从二进制"):
            msg_codec.encrypt_pub(b"test")

    def test_decrypt_pri_raises(self):
        with pytest.raises(NotImplementedError, match="未反推|未从二进制"):
            msg_codec.decrypt_pri(b"\x00" * 256)


# ============== 集成测试: 模拟完整抓包场景 ==============

class TestIntegration:
    """端到端: 模拟一次 client→server→client 的完整协议流"""

    def test_create_role_full_flow(self):
        """模拟 §11.4 创建角色完整调用链 (创角 CS101 → 成功 SC521)"""
        from msg_codec import send, receive, MSG_CS_CREATE_ROLE, MSG_SC_CREATE_ROLE_OK

        # 1. client 创角上行
        body_cs = b'{"account":"z1098210167","chr":"\xe8\x90\x8c\xe6\x96\xb0","job":0,"sex":1,"server_name":"\xe6\xaf\x94\xe5\xa5\x87","sdkid":"37504","secsdk":"0"}'
        cs_wire = send(MSG_CS_CREATE_ROLE, 900000, 12345, 37504, 0, body_cs)

        # 2. server 回包 (手动构造 SC521 成功响应)
        body_sc = b"\xe8\x90\x8c\xe6\x96\xb0/12345"  # 名字/roleID
        sc_wire = send(MSG_SC_CREATE_ROLE_OK, 0, 12345, 0, 0, body_sc)

        # 3. client 端解析 (抓包视角)
        cs_frames, cs_rest = receive(cs_wire)
        sc_frames, sc_rest = receive(sc_wire)
        assert cs_rest == sc_rest == b""
        assert cs_frames[0].msg_id == MSG_CS_CREATE_ROLE
        assert sc_frames[0].msg_id == MSG_SC_CREATE_ROLE_OK
        assert sc_frames[0].body == body_sc

    def test_tcp_fragmented_recv(self):
        """模拟 TCP 跨两次 recv 的粘包/半包"""
        body = b"x" * 200
        wire = msg_codec.send(529, 0, 0, 0, 0, body, 0)

        # 第一次 recv 拿到前 100 字节
        buf1 = wire[:100]
        frames, rest = msg_codec.receive(buf1)
        assert frames == []
        assert rest == buf1  # 不够一帧, 全部留存

        # 第二次 recv 拿到剩余
        buf2 = wire[100:]
        frames, rest = msg_codec.receive(rest + buf2)
        assert rest == b""
        assert len(frames) == 1
        assert frames[0].body == body

    def test_batch_message_concat(self):
        """批量包 (-1) 多个 msg 拼接在一起"""
        msgs = [
            frame_mod.Frame(20, 0, 0, 0, 0, 0, b"msg1_payload"),
            frame_mod.Frame(101, 900000, 0, 0, 0, 0, b"msg2_payload_longer"),
            frame_mod.Frame(529, 0, 0, 0, 0, 0, b"msg3"),
        ]
        # 模拟批量包: 把多个 frame body 串起来, 外层 -1 帧
        concat_body = b"".join(m.pack() for m in msgs)
        batch_wire = frame_mod.pack_frame(-1, 0, 0, 0, 0, 0, concat_body)

        # 解外层: 得到批量包 body
        frames, _ = msg_codec.receive(batch_wire)
        assert len(frames) == 1
        assert frames[0].msg_id == -1
        assert len(frames[0].body) == len(concat_body)


if __name__ == "__main__":
    sys.exit(pytest.main([__file__, "-v", "--tb=short"]))