from Crypto.Cipher import AES
import base64
import hashlib

class HAES:
    def __init__(self, key: bytes=b'Sixteen byte key') -> None:
        '''
        初始化需要指定密钥 key, key 必须是 16 24 32 字节的
        '''
        if isinstance(key, str):
            key = key.encode("utf8")
        if not key:
            key = b"_hon233_" * 4
        while len(key) < 32:
            key += key
        key = key[:32]
        self.cipher = AES.new(key, AES.MODE_ECB)

    def encode(self, context, rstr: bool=False, rbytes: bool=False):
        '''
        将传入的明文字符串加密

        :param context: 明文数据
        :param rstr: 如果此参数为 True 则返回解码后的字符串密文
        :param rbytes: 如果此参数为 True 则返回解码后的字节密文
        '''
        if rstr and rbytes:
            raise ValueError("不能同时指定输出字符串和字节")
        if isinstance(context, str):
            context = context.encode('utf8')
            if not rbytes:
                rstr = True
        while len(context) % AES.block_size:
            context += b'\x00'
        eData = self.cipher.encrypt(context)
        if rstr:
            return base64.encodebytes(eData).decode('utf-8').replace('\n', '')
        return eData

    def decode(self, context, rstr: bool=False, rbytes: bool=False):
        '''
        将传入的密文还原成明文

        :param  context: str | bytes 密文数据
        :param  rstr: bool 如果此参数为 True 则返回解码后的字符串明文
        :param rbytes: 如果此参数为 True 则返回解码后的字节明文
        '''
        if rstr and rbytes:
            raise ValueError("不能同时指定输出字符串和字节")
        if isinstance(context, str):
            eData = base64.decodebytes(context.encode('utf-8'))
            if not rbytes:
                rstr = True
        else:
            eData = context
        dData = self.cipher.decrypt(eData).rstrip(b'\x00')
        if rstr:
            return dData.decode('utf-8').replace('\n', '')
        return dData


if __name__ == "__main__":
    key = "_hon233_"
    obj = HAES(key)
    tmp = "需要加密的字符串"
    tmp2 = obj.encode(tmp)
    print(tmp2)
    tmp3 = obj.decode(tmp2)
    print(tmp3)