file_tools.py 2.9 KB
import os
import re
import zipfiles

import rarfile
# from zipfile import ZipFile
from pyzipper import AESZipFile


def file_write(file, file_path):
    with open(file_path, 'wb+') as f:
        for chunk in file.chunks():
            f.write(chunk)


def write_zip_file(dir_name, zipfile_path):
    if not os.path.isdir(dir_name):
        return
    with zipfiles.ZipFile(zipfile_path, 'w') as z:
        for root, dirs, files in os.walk(dir_name):
            root_target_path = root.replace(dir_name, '')
            for single_file in files:
                src_file_path = os.path.join(root, single_file)
                file_target_path = os.path.join(root_target_path, single_file)
                z.write(src_file_path, file_target_path)


def get_pwd_list_from_str(doc_name):
    try:
        pwd_list = re.findall(r'\d{6}', doc_name)
        return pwd_list
    except Exception as e:
        return []


def extract_zip_or_rar(file_path, extract_path, pwd_list=[]):
    if file_path.endswith('.zip') or file_path.endswith('.ZIP'):
        if len(pwd_list) > 0:
            for password in pwd_list:
                try:
                    with AESZipFile(file_path) as zf:
                        zf.extractall(extract_path, pwd=bytes(password, 'utf-8'))
                except Exception as e:
                    continue
                else:
                    return True
            else:
                return False
        else:
            try:
                with zipfiles.ZipFile(file_path) as zf:
                    zf.extractall(extract_path)
            except Exception as e:
                return False
            else:
                return True
    elif file_path.endswith('.rar') or file_path.endswith('.RAR'):
        if len(pwd_list) > 0:
            for password in pwd_list:
                try:
                    with rarfile.RarFile(file_path) as rf:
                        rf.extractall(extract_path, pwd=password)
                except Exception as e:
                    continue
                else:
                    return True
            else:
                return False
        else:
            try:
                with rarfile.RarFile(file_path) as rf:
                    rf.extractall(extract_path)
            except Exception as e:
                return False
            else:
                return True
    else:
        return False


def get_file_paths(input_path, suffix_list):
    """

    Args:
        input_path: str 目标目录
        suffix_list: list 搜索的文件的后缀列表

    Returns: list 搜索到的相关文件绝对路径列表

    """
    for parent, _, filenames in os.walk(input_path):
        for filename in filenames:
            for suffix in suffix_list:
                if filename.endswith(suffix):
                    file_path = os.path.join(parent, filename)
                    break
            else:
                continue
            yield file_path