SM4 加扰认证加密实战：Encrypt-then-MAC 实现与生产陷阱

from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes

def derive_keys(master_key: bytes, salt: bytes, info: bytes) -> tuple[bytes, bytes]:
    """从主密钥派生加密密钥和 MAC 密钥"""
    hkdf = HKDF(
        algorithm=hashes.SM3(),
        length=48,       # 16 + 32
        salt=salt,
        info=info,
    )
    expanded = hkdf.derive(master_key)
    enc_key = expanded[:16]   # SM4-128
    mac_key = expanded[16:]   # HMAC-SM3
    return enc_key, mac_key

"""
SM4-Encrypt-then-MAC 认证加密实现
依赖: pip install gmssl cryptography
测试环境: Python 3.10+, gmssl 3.2.2, cryptography 42.0+
"""
import os
import struct
from typing import tuple

from gmssl import sm4, func
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.hmac import HMAC
from cryptography.hazmat.backends import default_backend
from hmac import compare_digest

# ---------- 参数 ----------
SM4_KEY_LEN   = 16   # SM4-128
SM4_BLOCK_LEN = 16   # 分组长度
IV_LEN        = 16   # SM4-CBC 固定 IV 长度
TAG_LEN       = 32   # HMAC-SM3 输出长度
HKDF_LEN      = 48   # enc_key(16) + mac_key(32)


def derive_keys(master_key: bytes, salt: bytes, context: bytes) -> tuple[bytes, bytes]:
    """HKDF-SM3 派生加密密钥和 MAC 密钥"""
    hkdf = HKDF(
        algorithm=hashes.SM3(),
        length=HKDF_LEN,
        salt=salt,
        info=context,
    )
    okm = hkdf.derive(master_key)
    return okm[:SM4_KEY_LEN], okm[SM4_KEY_LEN:]


def _pkcs7_pad(data: bytes, block_size: int = SM4_BLOCK_LEN) -> bytes:
    pad_len = block_size - (len(data) % block_size)
    return data + bytes([pad_len] * pad_len)


def _pkcs7_unpad(data: bytes, block_size: int = SM4_BLOCK_LEN) -> bytes:
    if len(data) == 0 or len(data) % block_size != 0:
        raise ValueError("invalid padding length")
    pad_len = data[-1]
    if pad_len < 1 or pad_len > block_size:
        raise ValueError("invalid padding byte")
    if data[-pad_len:] != bytes([pad_len] * pad_len):
        raise ValueError("padding mismatch")
    return data[:-pad_len]


def encrypt(plaintext: bytes, master_key: bytes, context: bytes) -> bytes:
    """
    Encrypt-then-MAC：先 SM4-CBC 加密，再 HMAC-SM3 签名。
    输出格式：iv(16) || ciphertext || tag(32)
    """
    # 1. 派生密钥（salt 必须唯一，context 是应用标识）
    salt = os.urandom(16)
    enc_key, mac_key = derive_keys(master_key, salt, context)

    # 2. 生成随机 IV
    iv = os.urandom(SM4_BLOCK_LEN)

    # 3. SM4-CBC 加密（PKCS7 填充）
    padded = _pkcs7_pad(plaintext)
    cipher = Cipher(algorithms.SM4(enc_key), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()
    ciphertext = encryptor.update(padded) + encryptor.finalize()

    # 4. HMAC-SM3 对 iv || ciphertext 计算标签
    h = HMAC(mac_key, hashes.SM3(), backend=default_backend())
    h.update(iv)
    h.update(ciphertext)
    tag = h.finalize()

    # 5. 打包输出
    return salt + iv + ciphertext + tag


def decrypt(ciphertext_with_tag: bytes, master_key: bytes, context: bytes) -> bytes:
    """
    Decrypt-then-verify（EtM 反序）：先验 HMAC，再解密。
    """
    # 1. 解析结构
    offset = 0
    salt = ciphertext_with_tag[offset:offset+16];   offset += 16
    iv   = ciphertext_with_tag[offset:offset+16];   offset += 16
    tag  = ciphertext_with_tag[-32:];                 offset = len(ciphertext_with_tag) - 32
    ciphertext = ciphertext_with_tag[32:-32]

    # 2. 派生密钥（用同一 salt/context，确定性）
    enc_key, mac_key = derive_keys(master_key, salt, context)

    # 3. 先验证 MAC（失败直接抛异常，不泄露解密结果）
    h = HMAC(mac_key, hashes.SM3(), backend=default_backend())
    h.update(iv)
    h.update(ciphertext)
    expected_tag = h.finalize()
    if not compare_digest(expected_tag, tag):
        raise ValueError("integrity check failed: tag mismatch")

    # 4. SM4-CBC 解密
    cipher = Cipher(algorithms.SM4(enc_key), modes.CBC(iv), backend=default_backend())
    decryptor = cipher.decryptor()
    padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize()

    # 5. 去填充
    return _pkcs7_unpad(padded_plaintext)


# ---------- 测试向量 ----------
if __name__ == "__main__":
    master_key = os.urandom(32)   # 32 字节主密钥
    context = b"gm.okpki.com:field-encryption:v1"

    # 测试 1：基本加解密
    plaintext = b"Hello, SM4-Encrypt-then-MAC!"
    ct = encrypt(plaintext, master_key, context)
    pt = decrypt(ct, master_key, context)
    assert pt == plaintext, f"round-trip failed: {pt!r} != {plaintext!r}"
    print(f"[PASS] round-trip: {pt.decode()}")

    # 测试 2：篡改检测
    tampered = bytearray(ct)
    tampered[20] ^= 0xFF           # 翻转密文第 5 个字节
    try:
        decrypt(bytes(tampered), master_key, context)
        print("[FAIL] tampering not detected!")
    except ValueError as e:
        print(f"[PASS] tamper detected: {e}")

    # 测试 3：不同 context 密钥不兼容
    ct2 = encrypt(plaintext, master_key, b"different-context")
    try:
        decrypt(ct2, master_key, context)
        print("[FAIL] cross-context not detected!")
    except ValueError:
        print("[PASS] cross-context rejected")

    # 测试 4：空字符串
    ct0 = encrypt(b"", master_key, context)
    pt0 = decrypt(ct0, master_key, context)
    assert pt0 == b""
    print(f"[PASS] empty string: len={len(ct0)} bytes (iv=16 + ct=16[pad] + tag=32)")

    # 测试 5：大数据量（1MB）
    big = os.urandom(1024 * 1024)
    ct_big = encrypt(big, master_key, context)
    pt_big = decrypt(ct_big, master_key, context)
    assert pt_big == big
    print(f"[PASS] 1MB round-trip: input={len(big)}B output={len(ct_big)}B")

    print("\n所有测试通过。")