import os

from gmssl import sm4

def encrypt_ecb(data, key):
    cipher = sm4.CryptSM4()
    cipher.set_key(key.encode('utf-8'), sm4.SM4_ENCRYPT)
    encrypted_data = cipher.crypt_ecb(data.encode('utf-8'))
    return encrypted_data.hex()

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

def encrypt_cbc_with_iv(data, key):
    cipher = sm4.CryptSM4()
    cipher.set_key(key.encode('utf-8'), sm4.SM4_ENCRYPT)

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

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

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

def decrypt_cbc_with_iv(encrypted_bytes, key):
    cipher = sm4.CryptSM4()
    cipher.set_key(key.encode('utf-8'), sm4.SM4_DECRYPT)

    # 提取IV
    iv = encrypted_bytes[:16]

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

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

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