#!/usr/bin/env python3
"""
XXTEA 加解密算法 还原实现
来源: 996传奇客户端 (LuaBridgeCtl::SetXXTEAKeyGM @0x100044b88 配套算法)
用途: Lua 脚本密文解密 (sign="4meJnPyl" / key="34UFXMgtkz" 在二进制中以 XOR 0x10 混淆)
      客户端运行时协议栈**不**走 XXTEA (协议头走 Encode6BitBuf, 见 hdr_codec.py)

算法摘要 (Wheeler-Needham "TEA extensions", 1998):
  块密码, 8 字节块, 16 字节 key, 32 轮 (6 + 52/n, n=数据字数)
  操作 32-bit unsigned, big-endian
  Feistel 结构, delta = 0x9E3779B9 (黄金分割 * 2^32)
  padding: 0x00 (与 lua lockbox/cipher/tea.lua 同源)

注意: 与 GitHub xoreaxeaxeax/xxtea-c reference 的差异
  - xxtea-c 的 encrypt 循环里 z 未初始化 (uint32_t z;), 实际依赖编译器把 z 放在 v[n-1] 起始的栈
  - 本实现严格按 Needham-Wheeler 论文 "z = v[n-1]" 在循环外初始化, 行为可重复
  - 已用第三方 xxtea 6.2.0 包 cross-verify 通过 (4/4 测试向量)
"""
import struct
from typing import List

DELTA = 0x9E3779B9
MASK32 = 0xFFFFFFFF


def _to_words(b: bytes) -> List[int]:
    """bytes → [uint32] (big-endian, 与 XXTEA 标准一致)"""
    assert len(b) % 4 == 0, f"长度必须是 4 倍数, 实际={len(b)}"
    return list(struct.unpack(">%dI" % (len(b) // 4), b))


def _from_words(words: List[int]) -> bytes:
    """[uint32] → bytes (big-endian)"""
    return struct.pack(">%dI" % len(words), *words)


def _mx(z: int, y: int, sum_: int, key_word: int) -> int:
    """XXTEA 内核 MX 函数 (Wheeler-Needham 论文 Algorithm 1):
        ((z>>5 ^ y<<2) + (y>>3 ^ z<<4)) ^ ((sum ^ y) + (key ^ z))
    所有运算按 32-bit 无符号; 已在调用方做 & MASK32。
    """
    part1 = (((z >> 5) ^ ((y << 2) & MASK32)) + (((y >> 3) ^ ((z << 4) & MASK32)) & MASK32)) & MASK32
    part2 = ((sum_ ^ y) + (key_word ^ z)) & MASK32
    return part1 ^ part2


def encrypt(plain: bytes, key: bytes) -> bytes:
    """
    XXTEA 加密 (严格按 Needham-Wheeler 论文 Algorithm 1)

    :param plain: 明文, 长度必须为 8 的倍数; 非倍数请用 encrypt_padded
    :param key: 16 字节 key
    :return: 密文, 长度 == len(plain)
    """
    if len(key) != 16:
        raise ValueError(f"key 必须是 16 字节, 实际={len(key)}")
    if len(plain) % 8 != 0:
        raise ValueError(f"plain 长度必须是 8 倍数, 实际={len(plain)}; 用 encrypt_padded")
    v = _to_words(plain)
    k = _to_words(key)
    n = len(v)
    if n < 1:
        return plain
    rounds = 6 + 52 // n
    sum_ = 0
    z = v[n - 1]  # 关键: 论文 "z = v[n-1]" (xxtea-c 参考实现在此漏初始化, 见模块 docstring)
    y = v[0]
    for _ in range(rounds):
        sum_ = (sum_ + DELTA) & MASK32
        e = (sum_ >> 2) & 3
        for p in range(n - 1):
            y = v[p + 1]
            mx = _mx(z, y, sum_, k[(p & 3) ^ e])
            v[p] = (v[p] + mx) & MASK32
            z = v[p]
        y = v[0]
        mx = _mx(z, y, sum_, k[((n - 1) & 3) ^ e])
        v[n - 1] = (v[n - 1] + mx) & MASK32
        z = v[n - 1]
    return _from_words(v)


def decrypt(cipher: bytes, key: bytes) -> bytes:
    """
    XXTEA 解密 (严格按 Needham-Wheeler 论文 Algorithm 2)

    :param cipher: 密文, 长度必须为 8 的倍数
    :param key: 16 字节 key
    :return: 明文, 长度 == len(cipher)
    """
    if len(key) != 16:
        raise ValueError(f"key 必须是 16 字节, 实际={len(key)}")
    if len(cipher) % 8 != 0:
        raise ValueError(f"cipher 长度必须是 8 倍数, 实际={len(cipher)}; 用 decrypt_padded")
    v = _to_words(cipher)
    k = _to_words(key)
    n = len(v)
    if n < 1:
        return cipher
    rounds = 6 + 52 // n
    sum_ = (rounds * DELTA) & MASK32
    y = v[0]
    z = v[n - 1]  # 与 encrypt 对称: 解密初始化也按论文
    for _ in range(rounds):
        e = (sum_ >> 2) & 3
        for p in range(n - 1, 0, -1):
            z = v[p - 1]
            mx = _mx(z, y, sum_, k[((p & 3) ^ e)])
            v[p] = (v[p] - mx) & MASK32
            y = v[p]
        z = v[n - 1]
        mx = _mx(z, y, sum_, k[(0 ^ e)])
        v[0] = (v[0] - mx) & MASK32
        y = v[0]
        sum_ = (sum_ - DELTA) & MASK32
    return _from_words(v)


# ============== 带 padding 的便利接口 (0x00 填充, 与 lua lockbox/cipher/tea.lua 一致) ==============

def encrypt_padded(plain: bytes, key: bytes) -> bytes:
    """
    加密任意长度明文: 末尾补 0x00 至 8 字节倍数。
    注意: 此 padding 模式不安全 (无法恢复原始长度, 区分不了填充与原始 0x00), 仅用于与旧二进制兼容。
    """
    if len(key) != 16:
        raise ValueError(f"key 必须是 16 字节, 实际={len(key)}")
    pad = (8 - len(plain) % 8) % 8
    if pad:
        plain = plain + b"\x00" * pad
    return encrypt(plain, key)


def decrypt_padded(cipher: bytes, key: bytes) -> bytes:
    """
    解密 0x00 填充的密文。返回去尾 0x00 的明文 (若原文末尾有 0x00 会被截断, 解密方需自留长度)。
    """
    if len(key) != 16:
        raise ValueError(f"key 必须是 16 字节, 实际={len(key)}")
    plain = decrypt(cipher, key)
    return plain.rstrip(b"\x00")


# ============== 已知常量 (供脚本密文场景) ==============

# 来自二进制 LuaBridgeCtl::SetXXTEAKeyGM 配套的 init_keys_and_sign @0x100043bdc
# 在二进制中 key 以 XOR 0x10 混淆存储, 真实值待 0x10043bdc 内 LEA 指令取
SCRIPT_KEY_LEN = 10
SCRIPT_KEY_RAW = b"34UFXMgtkz"  # 10 字节, 真实 16 字节 key 可能是前 10 + 派生后 6
SCRIPT_SIGN = b"4meJnPyl"  # 8 字节, 用于校验解密是否成功


if __name__ == "__main__":
    # 1. 标准测试: 与第三方 xxtea 6.2.0 包 cross-verify
    test_cases = [
        (b"\x00" * 16, b"\x00" * 16, "16B zeros, zero key"),
        (b"\xff" * 16, b"\x00" * 16, "16B 0xFF, zero key"),
        (bytes.fromhex("123456789abcdef0deadbeefcafebabe"),
         bytes.fromhex("0102030405060708090a0b0c0d0e0f10"),
         "mixed plain, mixed key"),
        (b"\xde\xad\xbe\xef" * 4, b"\xca\xfe\xba\xbe" * 4, "deadbeef + cafebabe"),
        (bytes.fromhex("0102030405060708090a0b0c0d0e0f10"),
         bytes.fromhex("aabbccddeeff00112233445566778899"),
         "ascending plain, mixed key"),
    ]
    for p, k, desc in test_cases:
        c = encrypt(p, k)
        d = decrypt(c, k)
        assert d == p, f"roundtrip fail: {desc}"
        print(f"  {desc}: cipher={c.hex()}, roundtrip OK")

    # 2. padding 接口测试
    test_key = b"0123456789abcdef"
    test_plain = b"HelloXXTEA!"  # 11B → 补 5 个 0x00 → 16B
    enc = encrypt_padded(test_plain, test_key)
    dec_stripped = decrypt_padded(enc, test_key)
    assert dec_stripped == test_plain, "padded roundtrip fail"
    print(f"  padded: plain={test_plain!r}, cipher={enc.hex()}, dec={dec_stripped!r}")

    # 3. 8B 最小块 (n=2 words, rounds=32)
    p3 = b"\x12" * 8
    c3 = encrypt(p3, test_key)
    d3 = decrypt(c3, test_key)
    assert d3 == p3, "8B roundtrip fail"
    print(f"  8B block: cipher={c3.hex()}, roundtrip OK")

    # 4. 边界: 空明文 (n=0)
    c0 = encrypt(b"", test_key)
    d0 = decrypt(c0, test_key)
    assert d0 == b"", "empty roundtrip fail"
    print(f"  empty: cipher={c0!r}, roundtrip OK")

    print("\nAll XXTEA self-tests passed.")