# update: 2021-5-12-21
from Crypto.Cipher import AES
import base64
import hashlib
import zipfile
import pyminizip
import pyzipper
import os


class _AES(object):
    def __init__(self, key='7486E0264E999881D4EF7BEEDF05A9F7', model='ECB', iv='',
                 code_type='utf-8'):
        """
        code_type: utf-8/gbk
        """
        self.code_type = code_type
        self.model = {'ECB': AES.MODE_ECB, 'CBC': AES.MODE_CBC}[model]
        self.key = self.replenish(key)

        # --- create aes object ---
        if model == 'ECB':
            self.aes = AES.new(self.key, self.model)
        elif model == 'CBC':
            self.aes = AES.new(self.key, self.model, iv)

    def replenish(self, block, block_size=16):
        """block_size: AES key must be either 16, 24, or 32 bytes long"""
        block = block.encode(self.code_type)
        while len(block) % block_size != 0:
            block += b'\x00'
        return block

    def encrypt(self, text):
        text = self.replenish(text)
        encrypt_text = self.aes.encrypt(text)
        return base64.encodebytes(encrypt_text).decode().strip()

    def decrypt(self, text):
        text = base64.decodebytes(text.encode(self.code_type))
        decrypt_text = self.aes.decrypt(text)
        return decrypt_text.decode(self.code_type).strip('\0')


def text_to_b64(string):
    """b64编码"""
    return str(base64.b64encode(str(string).encode('utf-8')), 'utf-8')


def b64_to_text(string):
    """b64解码"""
    return str(base64.b64decode(str(string).encode('utf-8')), 'utf-8')


def text_to_aes(text, key='YourPassword'):
    """aes加密"""
    aes = _AES(key=key)
    return text_to_b64(aes.encrypt(text))


def aes_to_text(text, key='YourPassword'):
    """aes解密"""
    aes = _AES(key=key)
    return aes.decrypt(b64_to_text(text))


def get_big_file_md5(file_path, block_size=8 * 1024):
    """获取文件md5（默认使用8KB作为分块大小）"""
    m = hashlib.md5()
    with open(file_path, 'rb') as f:
        while True:
            block = f.read(block_size)
            if not block:
                break
            m.update(block)
    return m.hexdigest()


def file_to_zip(file_path):
    """生成zip压缩文件"""

    # --- encrypt file name --- todo 如果aes后超过255的长度就不做aes了
    file_name = file_path.split('/')[-1]
    save_name = text_to_b64(file_name)

    # --- define zip name ---
    file_md5 = get_big_file_md5(file_path)
    zip_name = f"{file_md5}.zip"
    dir_path = '/'.join(file_path.split('/')[:-1])
    zip_path = f"{dir_path}/{zip_name}"

    # --- writing ---
    with zipfile.ZipFile(zip_path, 'w', compression=zipfile.ZIP_DEFLATED, allowZip64=True) as f1:
        f1.write(file_path, arcname=save_name)


def zip_to_file():
    pass


def file_to_zip_v2(file_path, key='<YourPassword@2021>', compress_level=5):
    """
    文件转zip
    compress_level(int) between 1 to 9, 1 (more fast) <---> 9 (more compress) or 0 (default)
    """
    # --- check parameter ---
    if not os.path.isfile(file_path):
        return dict(code=1, details='parameter error!')

    # --- encrypt file name ---
    file_name = file_path.split('/')[-1]
    dir_path = '/'.join(file_path.split('/')[:-1])
    save_name = text_to_aes(file_name, key=key)

    # --- check again ---
    is_encrypted = len(file_name) == 37 and file_name[32] == '.'
    if is_encrypted:
        return dict(code=2, details='parameter error!')

    # --- check system support ---
    if len(save_name) > 250:
        return dict(code=3, details='parameter error!')

    # --- rename file ---
    save_path = f"{dir_path}/{save_name}"
    os.rename(file_path, save_path)

    # --- define zip name ---
    file_md5 = get_big_file_md5(save_path)
    zip_path = f"{dir_path}/{file_md5}.{save_name[:4]}"

    try:
        # --- writing ---
        """
        Args:
            1. src file path (string)
            2. src file prefix path (string) or None (path to prepend to file)
            3. dst file path (string)
            4. password (string) or None (to create no-password zip)
            5. compress_level(int) between 1 to 9, 1 (more fast) <---> 9 (more compress) or 0 (default)
        """
        pyminizip.compress(save_path, None, zip_path, key, compress_level)

        # --- remove file ---
        os.remove(save_path)
        return dict(code=0, details='ended.')

    except Exception as exception:

        # --- revert ---
        os.rename(save_path, file_path)

        import traceback
        print(traceback.format_exc())
        return dict(code=-1, details=f"{traceback.format_exc()}")


def zip_to_file_v2(file_path, key='<YourPassword@2021>', no_path=True):
    """zip转file"""

    # --- check parameter ---
    if not os.path.isfile(file_path):
        return dict(code=1, details='parameter error!')

    # --- check again ---
    file_name = file_path.split('/')[-1]
    is_encrypted = len(file_name) == 37 and file_name[32] == '.'
    if not is_encrypted:
        return dict(code=2, details='parameter error!')

    # --- record ---
    dir_path = '/'.join(file_path.split('/')[:-1])
    names = os.listdir(dir_path)
    before_files = set(f"{dir_path}/{name}" for name in names if not os.path.isdir(f"{dir_path}/{name}"))

    try:
        # --- writing ---
        """
        Args:
            1. src file path (string)
            2. password (string) or None (to unzip encrypted archives)
            3. dir path to extract files or None (to extract in a specific dir or cwd)
            4. withoutpath (exclude path of extracted)
        """
        pyminizip.uncompress(file_path, key, dir_path, no_path)

        # --- record ---
        names = os.listdir(dir_path)
        after_files = set(f"{dir_path}/{name}" for name in names if not os.path.isdir(f"{dir_path}/{name}"))

        # --- rename ---
        for path in list(after_files - before_files):
            name = path.split('/')[-1]
            raw_name = aes_to_text(name, key=key)
            os.rename(path, f"{dir_path}/{raw_name}")

        # --- remove file ---
        os.remove(file_path)
        return dict(code=0, details='ended.')

    except Exception as exception:

        # --- revert ---
        names = os.listdir(dir_path)
        after_files = set(f"{dir_path}/{name}" for name in names if not os.path.isdir(f"{dir_path}/{name}"))
        for path in list(after_files - before_files):
            os.remove(path)

        import traceback
        print(traceback.format_exc())
        return dict(code=-1, details=f"{traceback.format_exc()}")


def file_to_zip_v2_1(file_path, key='<YourPassword@2021>', compress_level=1):
    """file转zip（双文件保存版，其中一个文件保存加密文件名）"""

    # --- check parameter ---
    if not os.path.isfile(file_path):
        return dict(code=1, details='parameter error!')

    # --- check again ---
    file_name = file_path.split('/')[-1]
    is_encrypted = len(file_name) == 37 and file_name[32] == '.'
    if is_encrypted:
        return dict(code=2, details='parameter error!')

    # --- create name file ---
    dir_path = '/'.join(file_path.split('/')[:-1])
    file_name_aes = text_to_aes(file_name, key=key)
    name_file_path = f"{dir_path}/{file_name_aes[:4]}"
    with open(name_file_path, 'w') as f1:
        f1.write(file_name_aes)

    # --- rename file ---
    file_md5 = get_big_file_md5(file_path)
    data_file_path = f"{dir_path}/{file_md5}"
    os.rename(file_path, data_file_path)

    # --- define zip name ---
    zip_path = f"{dir_path}/{file_md5}.{file_name_aes[:4]}"

    try:
        # --- writing ---
        # pyminizip.compress(save_path, None, zip_path, key, compress_level)
        """
        Args:
            1. src file LIST path (list)
            2. src file LIST prefix path (list) or []
            3. dst file path (string)
            4. password (string) or None (to create no-password zip)
            5. compress_level(int) between 1 to 9, 1 (more fast) <---> 9 (more compress)
            6. optional function to be called during processing which takes one argument, the count of how many files have been compressed
        """
        pyminizip.compress_multiple([data_file_path, name_file_path], ['', ''], zip_path, key, compress_level)

        # --- remove file ---
        os.remove(data_file_path)
        os.remove(name_file_path)
        return dict(code=0, details='ended.')

    except Exception as exception:

        # --- revert ---
        os.remove(name_file_path)
        os.rename(data_file_path, file_path)

        import traceback
        print(traceback.format_exc())
        return dict(code=-1, details=f"{traceback.format_exc()}")


def zip_to_file_v2_1(file_path, key='<YourPassword>', no_path=True):
    """zip转file（双文件保存版，其中一个文件保存加密文件名）"""

    # --- check parameter ---
    if not os.path.isfile(file_path):
        return dict(code=1, details='parameter error!')

    # --- check again ---
    file_name = file_path.split('/')[-1]
    is_encrypted = len(file_name) == 37 and file_name[32] == '.'
    if not is_encrypted:
        return dict(code=2, details='parameter error!')

    # --- record ---
    dir_path = '/'.join(file_path.split('/')[:-1])
    names = os.listdir(dir_path)
    before_files = set(f"{dir_path}/{name}" for name in names if not os.path.isdir(f"{dir_path}/{name}"))

    try:
        # --- writing ---
        """
        Args:
            1. src file path (string)
            2. password (string) or None (to unzip encrypted archives)
            3. dir path to extract files or None (to extract in a specific dir or cwd)
            4. withoutpath (exclude path of extracted)
        """
        pyminizip.uncompress(file_path, key, dir_path, no_path)

        # --- record ---
        names = os.listdir(dir_path)
        after_files = set(f"{dir_path}/{name}" for name in names if not os.path.isdir(f"{dir_path}/{name}"))

        # --- check ---
        temp_file = list(after_files - before_files)
        if len(temp_file) != 2:
            raise Exception('something is wrong!')

        # --- get file name ---
        raw_name = str()
        for path in temp_file:
            name = path.split('/')[-1]
            if len(name) == 4:
                with open(path, 'r') as f1:
                    raw_name = aes_to_text(f1.read(), key=key)
                os.remove(path)

        # --- rename ---
        for path in list(after_files - before_files):
            name = path.split('/')[-1]
            if len(name) == 32:
                os.rename(path, f"{dir_path}/{raw_name}")

        # --- remove file ---
        os.remove(file_path)
        return dict(code=0, details='ended.')

    except Exception as exception:

        # --- revert ---
        names = os.listdir(dir_path)
        after_files = set(f"{dir_path}/{name}" for name in names if not os.path.isdir(f"{dir_path}/{name}"))
        for path in list(after_files - before_files):
            os.remove(path)

        import traceback
        print(traceback.format_exc())
        return dict(code=-1, details=f"{traceback.format_exc()}")


def file_to_zip_v3(file_path, key='YourPassword', aes_bits=128):
    """
    生成aes加密zip压缩文件
    aes_bits: 128/192/256
    """
    # --- check parameter ---
    if not os.path.isfile(file_path):
        return dict(code=1, details='parameter error!')

    # --- encrypt file name ---
    file_name = file_path.split('/')[-1]
    save_name = text_to_aes(file_name, key=key)

    # --- check system support ---
    if len(save_name) > 250:
        return dict(code=2, details='parameter error!')

    # --- define zip name ---
    file_md5 = get_big_file_md5(file_path)
    zip_name = f"{file_md5}.zip"
    dir_path = '/'.join(file_path.split('/')[:-1])
    zip_path = f"{dir_path}/{zip_name}"

    # --- writing ---
    with pyzipper.AESZipFile(zip_path, 'w', compression=pyzipper.ZIP_LZMA) as f1:
        f1.setpassword(key.encode())
        f1.setencryption(pyzipper.WZ_AES, nbits=aes_bits)
        f1.write(file_path, arcname=save_name)
    return dict(code=0, details='ended.')


def zip_to_file_v3(file_path, key='YourPassword'):
    """生成aes加密zip压缩文件"""

    dir_path = '/'.join(file_path.split('/')[:-1])

    # --- reading ---
    with pyzipper.AESZipFile(file_path) as f1:
        f1.setpassword(key.encode())
        file_list = f1.namelist()
        for file_name in file_list:
            # --- decrypt file name ---
            raw_name = aes_to_text(file_name, key=key)

            # --- writing ---
            with open(f"{dir_path}/{raw_name}", 'wb') as f2:
                f2.write(f1.read(file_name))
    return dict(code=0, details='ended.')


def file_to_7z(file_path, key='YourPassword'):
    pass


def folder_to_zip(folder_path):
    pass


if __name__ == '__main__':
    # text1 = '0/588564d4-3bca-4ebf-83d8-1155963fbfa9/3066753951/1593578486/1595174400'
    text2 = 'eFVqYzZjZ3hneVZ0T09uVVlHTm1rUT09'
    # text2 = text_to_aes(text1, '123456')
    test3 = aes_to_text(text2, '123456')
    # print('密文:', text2)
    print('明文:', test3)