import base64
import os

from gmssl import sm4
"""
输入：
消息或者加密后消息：字符串类型
key：字节类型
输出：
消息或者加密后消息：base64的字符串类型
"""

def encrypt_ecb(data: str, key: bytes) -> str:
    cipher = sm4.CryptSM4()
    cipher.set_key(key, sm4.SM4_ENCRYPT)
    encrypted_data = cipher.crypt_ecb(data.encode('utf-8'))
    return base64.b64encode(encrypted_data).decode('utf-8')


def decrypt_ecb(encrypted_hex: str, key: bytes) -> str:
    cipher = sm4.CryptSM4()
    cipher.set_key(key, sm4.SM4_DECRYPT)
    decrypted_data = cipher.crypt_ecb(bytes.fromhex(encrypted_hex))
    return decrypted_data.decode('utf-8')


def encrypt_cbc_with_iv(data: str, key: bytes) -> str:
    cipher = sm4.CryptSM4()
    cipher.set_key(key, sm4.SM4_ENCRYPT)

    # 生成随机的16字节IV
    iv = os.urandom(16)

    # 加密数据
    encrypted_data = cipher.crypt_cbc(iv, data.encode('utf-8'))

    # 将IV和加密后的数据拼接在一起
    return base64.b64encode(iv + encrypted_data).decode('utf-8')


def decrypt_cbc_with_iv(encrypted_bytes_base64: str, key) -> str:
    encrypted_bytes = base64.b64decode(encrypted_bytes_base64)
    cipher = sm4.CryptSM4()
    cipher.set_key(key, sm4.SM4_DECRYPT)

    # 提取IV
    iv = encrypted_bytes[:16]

    # 提取加密后的数据
    encrypted_data = encrypted_bytes[16:]

    # 解密数据
    decrypted_data = cipher.crypt_cbc(iv, encrypted_data)
    return base64.b64decode(decrypted_data).decode('utf-8')


if __name__ == "__main__":
    # 示例数据和密钥
    data = "Hello, SM4!"
    key = b"1234567890abcdef"

    # 加密
    encrypted_data = encrypt_ecb(data, key)
    print(f"Encrypted: {encrypted_data}")

    # 解密
    decrypted_data = decrypt_ecb(encrypted_data, key)
    print(f"Decrypted: {decrypted_data}")

    # 示例数据和密钥
    data = "Hello, SM4 CBC with random IV!"
    key = b"1234567890abcdef"

    # 加密
    encrypted_data = encrypt_cbc_with_iv(data, key)
    print(f"Encrypted: {encrypted_data.hex()}")

    # 解密
    decrypted_data = decrypt_cbc_with_iv(encrypted_data, key)
    print(f"Decrypted: {decrypted_data}")