aead 模式，sm4，aes 的gcm模式，文件名：aead_alg ,bytes 操作

"""
Author: tanglei
DateTime:2024-12-27 完成
微信：ciss_cedar
欢迎一起学习
加密算法：["AES","SM4"] 的GCM模式，特别适合于接口中传送数据
加密，认证，完整性 一起完成
"""
from enum import Enum
from Cryptodome.Util.Padding import pad, unpad
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.backends.openssl import backend
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes


class CipherMode(Enum):
    GCM='GCM'

class AlgName(Enum):
    AES,SM4='AES', 'SM4'


class MyAeadCipher:
    backend = default_backend()
    initialization_vector=b'0102030405060708'
    block_size = 16
    min_tag_length=16
    style = 'pkcs7'
    MODE_NAME = ['GCM']
    ALG_NAME = ['AES', 'SM4']
    key_sizes = frozenset([128, 192, 256])

    def __init__(self,key,aad,tag,mode,alg_name):
           self.key=key
           self.aad=aad
           #my_variable if my_variable is not None else "默认值"
           self.tag= tag if tag is not None else b'01020304050607080102030405060708'
           self.mode=self.get_model(mode,tag,aad)
           self.alg_name=self.valid_alg_name(alg_name)
           #print(f'tag=',self.tag.hex())

    @staticmethod
    def get_model(mode, tag,aad):
        if mode not in MyAeadCipher.MODE_NAME:
            raise "unsupported modes, confirm the input Mode name"

        #tag = pad(tag, MyAeadCipher.block_size, MyAeadCipher.style)
        padded_aad = pad(aad, MyAeadCipher.block_size, MyAeadCipher.style)
        #mode = modes.GCM(padded_aad or b'')
        if tag is None:
            modes_dict = {
                'GCM': modes.GCM(padded_aad),
            }
        else:
            modes_dict = {
                'GCM': modes.GCM(padded_aad, tag)
            }
        return modes_dict.get(mode)

    @staticmethod
    def valid_alg_name(alg_name):
        if alg_name not in MyAeadCipher.ALG_NAME:
            raise "unsupported alg, confirm the input ALG_NAME name"
        return alg_name

    def encrypt_bytes(self, plaintext):
        """
        输入字符串，完成加密操作
        :param plaintext 明文字符串 bytes
        :return: 16进制密文字符串
        """
        # bytes后的填充数据
        data_padding = pad(plaintext, self.block_size, self.style)
        # 加密对象AES,SM4

        algorithm = getattr(algorithms, self.alg_name)

        my_alg = Cipher(algorithm(self.key), self.mode,backend)
        encryptor = my_alg.encryptor()
        # 加密操作完成
        encrypt_bytes = encryptor.update(data_padding) + encryptor.finalize()
        tag_bytes = encryptor.tag
        return encrypt_bytes,tag_bytes

    def decrypt_bytes(self, ciphertext):
        """
        输入16进制密文数据，返回原文
        :param ciphertext: 16进制密文数据字符串
        :return:
        """

        algorithm = getattr(algorithms, self.alg_name)
        my_alg = Cipher(algorithm(self.key), self.mode,backend)
        decrypt = my_alg.decryptor()
        # 解密操作完成
        pad_decrypt_value = decrypt.update(ciphertext) + decrypt.finalize()
        decrypt_bytes = unpad(pad_decrypt_value,  self.block_size, self.style)
        # 转为大写的16进制字符串输出
        return decrypt_bytes


def main():
    key128 = '2934412A66B7A186DC35DC40E926F9EE'
    key192 = '2934412A66B7A186DC35DC40E926F9EE2934412A66B7A186'
    key256 = '2934412A66B7A186DC35DC40E926F9EE2934412A66B7A186DC35DC40E926F9EE'
    tag = None
    key = bytes.fromhex(key128)
    aad=b''
    alg_name = AlgName.AES.value
    mode = CipherMode.GCM.value

    my_alg = MyAeadCipher(key,aad, tag, mode, alg_name)
    source = b'1'

    cipher, tag = my_alg.encrypt_bytes(source)

   # tag=tag.encode()
    # tag = None
    my_alg = MyAeadCipher(key, aad,tag, mode, alg_name)
    plain = my_alg.decrypt_bytes(cipher)
    print(f'alg_name={alg_name},mode={mode}')
    print(f'key={key},tag={tag}')  # .hex().upper()
    print(f'data=', source)

    print(f'cipher={cipher}')
    print(f'cipher.Hex: {cipher.hex().upper()}')
    print(f'plain={plain}')
    # print(f'plain.decode= {plain.decode()}')
    # AES_ECB:7F0AFB6D64764FAF13D73EDBA225F142
    # AES_CBC:3C6A796A00C52E41B1FDECF4206DDA1C
    # AES_CTR:5F30BCC30D7634F2C1D6411B674BC569
    # SM4_ECB:4B622D9C120C4FE245567C496684E2F0
    # SM4_CBC:42A3696426D3F0DCBDD7A58E17352BA3
    # SM4_CTR:7489C0A4178D449B912D2F78AF2F22A2


if __name__ == '__main__':
   main()