import os

from gmssl import sm4


def encrypt_ecb(data, key):
    cipher = sm4.CryptSM4()
    cipher.set_key(key, sm4.SM4_ENCRYPT)
    encrypted_data = cipher.crypt_ecb(data)
    return encrypted_data


def decrypt_ecb(encrypted_data, key):
    cipher = sm4.CryptSM4()
    cipher.set_key(key, sm4.SM4_DECRYPT)
    decrypted_data = cipher.crypt_ecb(encrypted_data)
    return decrypted_data


def encrypt_cbc_with_iv(data, key):
    cipher = sm4.CryptSM4()
    cipher.set_key(key, sm4.SM4_ENCRYPT)

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

    # 加密数据
    encrypted_data = cipher.crypt_cbc(iv, data)

    # 将IV和加密后的数据拼接在一起
    return iv + encrypted_data


def decrypt_cbc_with_iv(encrypted_data, key):
    cipher = sm4.CryptSM4()
    cipher.set_key(key, sm4.SM4_DECRYPT)

    # 提取IV
    iv = encrypted_data[:16]

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

    # 解密数据
    decrypted_data = cipher.crypt_cbc(iv, encrypted_data)
    return decrypted_data


if __name__ == "__main__":
    # 示例数据和密钥
    data = b"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 = b"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}")