#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@Version: 1.0
@Python Version:3.6.6
@Author: ludq1
@Email: ludq1@chinaunicom.cn
@date: 2023/04/07 11:40:00
@Description:
"""

import base64
from typing import Optional

from Crypto.Cipher import AES


class AesUtilsForMysql:
    r"""
    Mysql AES加密 utils类，不依赖任何其他应用基础类
    """

    def __init__(self, key: str = None):
        r"""
        初始化
        :param key
        """
        self.key: str = key
        if self.key:
            self.real_aes_key: Optional[bytes] = self.gen_aes_key_by_mysql_key(self.key)
        else:
            self.real_aes_key: Optional[bytes] = None

    def pkcs7padding_to_encrypt_content(self, to_encrypt_content: str) -> str:
        r"""
        明文使用PKCS7填充

        最终调用AES加密方法时，传入的是一个byte数组，要求是16的整数倍，因此需要对明文进行处理

        :param to_encrypt_content
        :return
        """
        bytes_length = len(bytes(to_encrypt_content, encoding='utf-8'))
        # tips：utf-8编码时，英文占1个byte，而中文占3个byte
        padding_times = 16 - bytes_length % 16
        # tips：chr(padding)看与其它语言的约定，有的会使用'\0'
        padding_text = chr(padding_times) * padding_times
        return to_encrypt_content + padding_text

    def pkcs7unpadding_decrypted_content(self, decrypted_content: str) -> str:
        r"""
        处理使用PKCS7填充过的数据

        :param decrypted_content
        :return
        """
        length = len(decrypted_content)
        unpadding = ord(decrypted_content[length - 1])
        return decrypted_content[0:length - unpadding]

    def gen_aes_key_by_mysql_key(self, key: str = None) -> bytes:
        r"""
        生成mysql函数中使用的key实际对应的aes_key
        :param key
        """
        real_key = key if key else self.key
        final_key = bytearray(16)
        for i, c in enumerate(real_key):
            final_key[i % 16] ^= ord(real_key[i])
        return bytes(final_key)

    def encrypt(self, to_encrypt_content: str, key: str = None) -> str:
        r"""
        mysql AES加密

        模式 AES.MODE_ECB

        填充pkcs7

        加密后使用 base64 编码

        :param to_encrypt_content
        :param key
        :return

        """
        real_aes_key = self.gen_aes_key_by_mysql_key(key) if key else self.real_aes_key
        cipher = AES.new(real_aes_key, AES.MODE_ECB)
        # 将内容编码
        to_encrypt_content = self.pkcs7padding_to_encrypt_content(to_encrypt_content)
        byte_content = to_encrypt_content.encode('utf-8')
        byte_encrypted = cipher.encrypt(byte_content)
        # base64 编码
        byte_base64 = base64.b64encode(byte_encrypted)
        return str(byte_base64, 'utf-8')

    def decrypt(self, to_decrypt_content: str, key: str = None):
        r"""
        mysql AES解密

        模式 AES.MODE_ECB

        去填充pkcs7

        解密前使用 base64 解码

        :param to_decrypt_content
        :param key
        :return

        """
        real_aes_key = self.gen_aes_key_by_mysql_key(key) if key else self.real_aes_key
        cipher = AES.new(real_aes_key, AES.MODE_ECB)
        # 先将内容解码
        byte_base64 = to_decrypt_content.encode('utf-8')
        byte_encrypted = base64.b64decode(byte_base64)
        # 解码
        result = cipher.decrypt(byte_encrypted)
        result = str(result, 'utf-8')
        # 去除填充内容
        result = self.pkcs7unpadding_decrypted_content(result)
        return result