'''
 author: lvsongke@oneniceapp.com
 data:2019/09/09
'''

import base64
from Crypto.Cipher import AES
import random


class CryptoNotOne(object):

    def pkcs7padding(self, text):
        """
        明文使用PKCS7填充
        最终调用AES加密方法时，传入的是一个byte数组，要求是16的整数倍，因此需要对明文进行处理
        :param text: 待加密内容(明文)
        :return:
        """
        bs = AES.block_size  # 16
        length = len(text)
        bytes_length = len(bytes(text, encoding='utf-8'))
        # tips：utf-8编码时，英文占1个byte，而中文占3个byte
        padding_size = length if (bytes_length == length) else bytes_length
        padding = bs - padding_size % bs
        # tips：chr(padding)看与其它语言的约定，有的会使用'\0'
        padding_text = chr(padding) * padding
        return text + padding_text

    def pkcs7unpadding(self, text):
        """
        处理使用PKCS7填充过的数据
        :param text: 解密后的字符串
        :return:
        """
        length = len(text)
        unpadding = ord(text[length - 1])
        return text[0:length - unpadding]

    def encrypt(self, key, content):
        """
        AES加密
        key,iv使用同一个
        模式cbc
        填充pkcs7
        :param key: 密钥
        :param content: 加密内容
        :return:
        """
        key_bytes = bytes(key, encoding='utf-8')
        # iv = key_bytes
        cipher = AES.new(key_bytes, AES.MODE_ECB)
        # 处理明文
        content_padding = self.pkcs7padding(content)
        # 加密
        encrypt_bytes = cipher.encrypt(bytes(content_padding, encoding='utf-8'))
        # 重新编码
        result = str(base64.b64encode(encrypt_bytes), encoding='utf-8')
        return result

    def decrypt(self, key, content):
        """
        AES解密
         key,iv使用同一个
        模式ecb
        去填充pkcs7
        :param key:
        :param content:
        :return:
        """
        key_bytes = bytes(key, encoding='utf-8')
        # iv = key_bytes
        cipher = AES.new(key_bytes, AES.MODE_ECB)
        # base64解码
        encrypt_bytes = base64.b64decode(content)
        # 解密
        decrypt_bytes = cipher.decrypt(encrypt_bytes)
        # 重新编码
        result = str(decrypt_bytes, encoding='utf-8')
        # 去除填充内容
        result = self.pkcs7unpadding(result)
        return result

    def get_key(self, n):
        """
        获取密钥 n 密钥长度
        :return:
        """
        c_length = int(n)
        source = 'ABCDEFGHJKMNPQRSTWXYZabcdefhijkmnprstwxyz2345678'
        # source = 'always love you '
        # length = len(source) - 1
        result = ''
        for i in range(c_length):
            result += random.choice(source)
        return result


if __name__ == '__main__':
    text = '遇见你之前我从未想过结婚，遇见你之后我只想结婚而且那个人只能是你'
    crypto = CryptoNotOne()
    aes_key = crypto.get_key(16)
    print('密钥:' + aes_key)
    print('明文:' + text)
    encrypt_en = crypto.encrypt(aes_key, text)
    print("密文:" + encrypt_en)
    result = crypto.decrypt(aes_key, encrypt_en)
    print(result)
