"""虎牙 Udb AES 加解密工具。

基于 AES-128-ECB 与 0 填充。
"""
from __future__ import annotations

from Crypto.Cipher import AES

MASTER_KEY_STR = b"HuyaUdb1928374650qwertyuiop"


def encrypt_block(key16: bytes, block16: bytes) -> bytes:
    return AES.new(key16, AES.MODE_ECB).encrypt(block16)


def decrypt_block(key16: bytes, block16: bytes) -> bytes:
    return AES.new(key16, AES.MODE_ECB).decrypt(block16)


def zero_pad(data: bytes) -> bytes:
    """UdbAESUtil::encrypt 填充：补 '\\0' 到 16 的倍数。"""
    if len(data) % 16 == 0:
        return data
    return data + b"\x00" * (16 - len(data) % 16)


def udb_encrypt(key16: bytes, plain: bytes) -> bytes:
    """AES-128-ECB 加密，带零填充。"""
    p = zero_pad(plain)
    return AES.new(key16, AES.MODE_ECB).encrypt(p)


def udb_decrypt(key16: bytes, ct: bytes) -> bytes:
    """AES-128-ECB 解密，去除尾部零填充。"""
    return AES.new(key16, AES.MODE_ECB).decrypt(ct).rstrip(b"\x00")