import logging
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.backends import default_backend

logger = logging.getLogger(__name__)


def encrypt(key, plain_data, mode='ecb', iv='00' * 16):  # key/plain_data/iv可以传入16进制字符串或者bytes
    if not all(isinstance(i, (bytes, str)) for i in (key, plain_data, iv)):
        raise TypeError("传入了非法的数据类型，传入的数据必须是16进制字符串或者bytes类型")
    try:
        plain_data = plain_data if isinstance(plain_data, bytes) else bytes.fromhex(plain_data)
        key = key if isinstance(key, bytes) else bytes.fromhex(key)
        iv = iv if isinstance(iv, bytes) else bytes.fromhex(iv)
    except ValueError:
        raise ValueError("不是合法的十六进制字符串，无法转换为bytes")
    mode = mode.lower()
    if mode == 'cbc':
        mode = modes.CBC(iv)
    elif mode == 'ecb':
        mode = modes.ECB()
    else:
        raise ValueError("不支持的模式,当前sm4只支持cbc/ecb模式")
    # 使用PKCS7进行填充
    pad = padding.PKCS7(algorithms.SM4.block_size).padder()
    padded_plain_data = pad.update(plain_data) + pad.finalize()
    # 创建SM4对象
    cipher = Cipher(algorithms.SM4(key), mode, backend=default_backend())
    encryptor = cipher.encryptor()  # 加密
    logger.debug(f"开始进行sm4加密计算,算法模式{mode.name},加密密钥:{key.hex()},加密数据长度:{len(plain_data) / 2}字节")
    cipher_data = encryptor.update(padded_plain_data)
    logger.debug(f"完成sm4加密计算")
    return cipher_data.hex()


def decrypt(key, cipher_data, mode='ecb', iv='00' * 16):  # key/cipher_data/iv可以传入16进制字符串或者bytes
    if not all(isinstance(i, (bytes, str)) for i in (key, cipher_data, iv)):
        raise TypeError("传入了非法的数据类型，传入的数据必须是16进制字符串或者bytes类型")
    try:
        cipher_data = cipher_data if isinstance(cipher_data, bytes) else bytes.fromhex(cipher_data)
        key = key if isinstance(key, bytes) else bytes.fromhex(key)
        iv = iv if isinstance(iv, bytes) else bytes.fromhex(iv)
    except ValueError:
        raise ValueError("不是合法的十六进制字符串，无法转换为bytes")
    mode = mode.lower()
    if mode == 'cbc':
        mode = modes.CBC(iv)
    elif mode == 'ecb':
        mode = modes.ECB()
    else:
        raise ValueError("不支持的模式,当前sm4只支持cbc/ecb模式")
    # 创建SM4对象
    cipher = Cipher(algorithms.SM4(key), mode, backend=default_backend())
    decryptor = cipher.decryptor()  # 解密
    logger.debug(f"开始进行sm4解密计算,算法模式{mode.name},解密密钥:{key.hex()}")
    padded_plain_data = decryptor.update(cipher_data)
    # 去除填充
    unpad = padding.PKCS7(algorithms.SM4.block_size).unpadder()
    plain_data = unpad.update(padded_plain_data) + unpad.finalize()
    logger.debug(f"完成sm4解密计算,解密后的数据长度:{len(plain_data)}字节")
    return plain_data.hex()
