"""
symmetric_alg.py 对bytes的操作
Author: tanglei
DateTime:2024-11-18 完成
微信：ciss_cedar
欢迎一起学习
需要IV的加密模式：CBC、OFB、CFB、CFB8
加密算法：["AES","Camellia","TripleDES","Blowfish","CAST5","ARC4","IDEA","SEED","ChaCha20","SM4"]
#ChaCha20 需要nonce ARC4 有漏洞。
"""
import os
from enum import Enum

from Cryptodome.Util.Padding import pad, unpad
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes


class CipherMode(Enum):
    ECB,CBC,CFB,OFB,CTR='ECB', 'CBC', 'CFB', 'OFB', 'CTR'

class AlgName(Enum):
    AES ,Camellia,TripleDES,Blowfish,CAST5, IDEA, SEED, SM4=\
        ('AES', 'Camellia', 'TripleDES', 'Blowfish', 'CAST5', 'IDEA',
         'SEED', 'SM4')


class MySymCipher:
    block_size = 16
    style = 'pkcs7'
    MODE_NAME = ['ECB', 'CBC', 'CFB', 'OFB', 'CTR']
    ALG_NAME = ['AES', 'Camellia', 'TripleDES', 'Blowfish', 'CAST5', 'IDEA', 'SEED', 'SM4']
    key_sizes = frozenset([128, 192, 256])


    def __init__(self,key,iv,mode,alg_name):
           self.key=key
           self.iv=iv  #认为 nonce
           self.mode=self.get_model(mode,iv)
           self.alg_name=self.valid_alg_name(alg_name)
        #algorithm=getattr(hashes, self.hash_name)
    @staticmethod
    def get_model(mode, iv):
        if mode not in MySymCipher.MODE_NAME:
            raise "unsupported modes, confirm the input Mode name"
        modes_dict = {
            'ECB': modes.ECB(),
            'CBC': modes.CBC(iv),
            'CFB': modes.CFB(iv),
            'OFB': modes.OFB(iv),
            'CTR': modes.CTR(iv),# 认为iv 是nonce
        }
        return modes_dict.get(mode.upper())

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

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

        algorithm = getattr(algorithms, self.alg_name)

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

    def decrypt_bytes(self, cipher_bytes):
        """
        输入16进制密文数据，返回原文
        :param ciphertext: 16进制密文数据字符串
        :return:
        """
        algorithm = getattr(algorithms, self.alg_name)
        my_alg = Cipher(algorithm(self.key), self.mode)
        decrypt = my_alg.decryptor()
        # 解密操作完成
        pad_decrypt_value = decrypt.update(cipher_bytes) + decrypt.finalize()
        decrypt_bytes = unpad(pad_decrypt_value, self.block_size, self.style)
        # 转为大写的16进制字符串输出
        return decrypt_bytes

    @staticmethod
    def gen_random_bytes(length=16):
        # 使用软件算法，可以调用系统函数或者使用secrets生成随机数再进行派生
        # 当前使用32字符，16字节的随机数也是可以的。
        random_bytes= os.urandom(length)
        return random_bytes
def main():
    key128 = '2934412A66B7A186DC35DC40E926F9EE'
    key192 = '2934412A66B7A186DC35DC40E926F9EE2934412A66B7A186'
    key256 = '2934412A66B7A186DC35DC40E926F9EE2934412A66B7A186DC35DC40E926F9EE'
    iv = '86CD720D75F4622DBE96078A3CD1076E'
    key = bytes.fromhex(key128)
    alg_name = AlgName.AES.value
    mode = CipherMode.ECB.value
    iv = bytes.fromhex(iv)
    my_alg = MySymCipher(key, iv, mode, alg_name)
    source = b'1'

    cipher = my_alg.encrypt_bytes(source)
    plain = my_alg.decrypt_bytes(cipher)
    print(f'alg_name={alg_name},mode={mode}')
    print(f'key={key},iv={iv}')
    print(f'source=', 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()