"""a_bogus 依赖的底层密码学原语：SM3、RC4、自定义 Base64。

这三样都是公开算法，与抖音无关，可以独立验证正确性：
  - SM3  见 GB/T 32905-2016，本文件底部带官方测试向量自检
  - RC4  标准流密码
  - 自定义 Base64 只是把标准码表换成抖音自己那套字符顺序
"""

from __future__ import annotations

from typing import Iterable, List, Sequence

# --------------------------------------------------------------------------- #
# SM3
# --------------------------------------------------------------------------- #

_IV = [
    0x7380166F, 0x4914B2B9, 0x172442D7, 0xDA8A0600,
    0xA96F30BC, 0x163138AA, 0xE38DEE4D, 0xB0FB0E4E,
]
_MASK = 0xFFFFFFFF


def _rotl(x: int, n: int) -> int:
    n &= 31
    return ((x << n) | (x >> (32 - n))) & _MASK


def _ff(j: int, x: int, y: int, z: int) -> int:
    return (x ^ y ^ z) if j < 16 else ((x & y) | (x & z) | (y & z))


def _gg(j: int, x: int, y: int, z: int) -> int:
    return (x ^ y ^ z) if j < 16 else ((x & y) | ((~x & _MASK) & z))


def _p0(x: int) -> int:
    return x ^ _rotl(x, 9) ^ _rotl(x, 17)


def _p1(x: int) -> int:
    return x ^ _rotl(x, 15) ^ _rotl(x, 23)


def _cf(v: Sequence[int], block: bytes) -> List[int]:
    """压缩函数 CF(V, B)。"""
    w = [int.from_bytes(block[i * 4:i * 4 + 4], "big") for i in range(16)]
    for j in range(16, 68):
        w.append(
            _p1(w[j - 16] ^ w[j - 9] ^ _rotl(w[j - 3], 15))
            ^ _rotl(w[j - 13], 7)
            ^ w[j - 6]
        )
    w1 = [w[j] ^ w[j + 4] for j in range(64)]

    a, b, c, d, e, f, g, h = v
    for j in range(64):
        t = 0x79CC4519 if j < 16 else 0x7A879D8A
        ss1 = _rotl((_rotl(a, 12) + e + _rotl(t, j)) & _MASK, 7)
        ss2 = ss1 ^ _rotl(a, 12)
        tt1 = (_ff(j, a, b, c) + d + ss2 + w1[j]) & _MASK
        tt2 = (_gg(j, e, f, g) + h + ss1 + w[j]) & _MASK
        d = c
        c = _rotl(b, 9)
        b = a
        a = tt1
        h = g
        g = _rotl(f, 19)
        f = e
        e = _p0(tt2)

    return [x ^ y for x, y in zip(v, (a, b, c, d, e, f, g, h))]


def sm3(data: bytes) -> bytes:
    """SM3 摘要，返回 32 字节。"""
    bit_len = len(data) * 8
    padded = data + b"\x80"
    padded += b"\x00" * ((56 - len(padded) % 64) % 64)
    padded += bit_len.to_bytes(8, "big")

    v = list(_IV)
    for i in range(0, len(padded), 64):
        v = _cf(v, padded[i:i + 64])
    return b"".join(x.to_bytes(4, "big") for x in v)


def sm3_to_list(data: bytes | Iterable[int]) -> List[int]:
    """SM3 摘要，返回 32 个 int（a_bogus 内部按字节数组传递）。"""
    if not isinstance(data, (bytes, bytearray)):
        data = bytes(bytearray(data))
    return list(sm3(bytes(data)))


# --------------------------------------------------------------------------- #
# RC4
# --------------------------------------------------------------------------- #

def rc4(key: bytes, data: bytes) -> bytes:
    """标准 RC4。加解密同一函数。"""
    s = list(range(256))
    j = 0
    klen = len(key)
    for i in range(256):
        j = (j + s[i] + key[i % klen]) & 0xFF
        s[i], s[j] = s[j], s[i]

    out = bytearray(len(data))
    i = j = 0
    for n, ch in enumerate(data):
        i = (i + 1) & 0xFF
        j = (j + s[i]) & 0xFF
        s[i], s[j] = s[j], s[i]
        out[n] = ch ^ s[(s[i] + s[j]) & 0xFF]
    return bytes(out)


# --------------------------------------------------------------------------- #
# 自定义码表 Base64
# --------------------------------------------------------------------------- #

STD_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="


def b64_encode(data: bytes, alphabet: str, pad: bool = True) -> str:
    """按给定 64(+1) 字符码表做 Base64 编码。

    alphabet 长度为 64 时不补位；为 65 时最后一位作为填充字符。
    """
    if len(alphabet) not in (64, 65):
        raise ValueError(f"码表长度必须是 64 或 65，收到 {len(alphabet)}")
    fill = alphabet[64] if len(alphabet) == 65 else ""

    out: List[str] = []
    for i in range(0, len(data), 3):
        chunk = data[i:i + 3]
        n = int.from_bytes(chunk + b"\x00" * (3 - len(chunk)), "big")
        quad = [alphabet[(n >> shift) & 0x3F] for shift in (18, 12, 6, 0)]
        keep = len(chunk) + 1  # 1字节->2字符, 2字节->3字符, 3字节->4字符
        out.extend(quad[:keep])
        if pad and fill:
            out.extend(fill * (4 - keep))
    return "".join(out)


# --------------------------------------------------------------------------- #
# 自检
# --------------------------------------------------------------------------- #

# GB/T 32905-2016 附录 A 给出的两组标准测试向量
_SM3_VECTORS = [
    (b"abc",
     "66c7f0f462eeedd9d1f2d46bdc10e4e24167c4875cf2f7a2297da02b8f4ba8e0"),
    (b"abcd" * 16,
     "debe9ff92275b8a138604889c18e5a4d6fdb70e5387e5765293dcba39c0c5732"),
]


def self_test() -> None:
    """跑通即说明 SM3/RC4/Base64 实现无误。"""
    for msg, want in _SM3_VECTORS:
        got = sm3(msg).hex()
        assert got == want, f"SM3({msg[:8]!r}...) = {got}, 期望 {want}"

    # RC4 官方测试向量：key="Key", plaintext="Plaintext" -> BBF316E8D940AF0AD3
    assert rc4(b"Key", b"Plaintext").hex().upper() == "BBF316E8D940AF0AD3"

    # 用标准码表时，自定义 Base64 应与 base64 模块完全一致
    import base64
    import os
    for n in range(1, 40):
        raw = os.urandom(n)
        assert b64_encode(raw, STD_ALPHABET) == base64.b64encode(raw).decode()

    print("crypto self-test OK: SM3 / RC4 / custom-base64 均通过")


if __name__ == "__main__":
    self_test()