python引入其他目录的类 python引入文件

python 批量加密压缩文件引入多进程及修复bug

• 0x00 写在前面
• 0x00 修改后的demo

0x00 写在前面

前面的demo：python 批量加密压缩文件中存在一些bug，本文修复了存在的bug，并引入多进程（为何是多进程而不是多线程：因为在操作中需要调用DOS命令去启动WinRAR进行压缩，要启用新的进程，所以不能用多线程）

我尝试了一下多线程，我要写的脚本的作用是：批量提取文件的hash，并把hash、文件大小等信息录入MySQL数据库，建立了数据库池、引入了多线程，但是当处理了20w个样本文件后，发现了问题，就是样本的信息都入库了，但是样本并没有被压缩，我才意识到了自己用了子线程去创建了进程，结果就是程序并没有报错，样本并没有压缩。所以应该用多进程，所以以后在运用时要注意使用线程还是进程，如果程序中没有创建新的进程或者子进程就可以使用多线程，否则就使用多进程。应当注意的是，多线程与多进程能让效率提高的不是一星半点儿，但是相应的会占用更多的计算机资源：CPU、内存等；所以要视情况，指定线程或者进程的数量，不要把自己的计算机跑死了。

0x00 修改后的demo

这里是一个简化的demo，不再进行数据库的操作了就直接进行压缩处理：

# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/3/27 16:58
# @Author  : H
# @File    : zipfilev3.py

import math
import os
import hashlib
import time

from multiprocessing import Process


def getfielpath(filepath, sub=[]):
    if os.path.isdir(filepath):
        # 如果绝对路径下的文件夹
        for i in os.listdir(filepath):  # i文件名
            path2 = os.path.join(filepath, i)  # 拼接绝对路径
            if os.path.isdir(path2):  # 判断如果是文件夹,调用本身
                getfielpath(path2, sub)
            else:
                sub.append(path2)
    elif os.path.isfile(filepath):
        # 如果绝对路径下的文件
        sub.append(filepath)
    else:
        print("File or path doesn\'t exit")
    return sub


"""获取文件的hash"""


def getHash(filepath):
    datas = {}  # 结果存为json，以便后需
    f = open(filepath, "rb")
    rb = f.read()
    data = {'MD5': hashlib.md5(rb).hexdigest(),
            'SHA1': hashlib.sha1(rb).hexdigest(),
            'SHA256': hashlib.sha256(rb).hexdigest(),
            'filesize': os.path.getsize(filepath)}
    f.close()
    """文件hash由完整路径构成词典"""
    datas[filepath] = data
    return datas


"""调用DOS压缩加密压缩"""


def zipFile(datas, WinRARpath, password, outputpath):
    for k, v in datas.items():
        filepath = k
        oldname = filepath.split('\\')[-1]
        md5 = v['MD5']
        sha1 = v['SHA1']
        sha256 = v['SHA256']
        filesize = v['filesize']
    """以文件的sha256重命名：若不重命名则当文件名中有空格或特殊符号时，启动DOS命令会失败"""
    # newfilename = filepath.replace(oldname, sha256)
    # repalce在文件名和路径名相同时会出错
    newfilename = ''
    for i in filepath.split("\\")[0:-1]:
        newfilename = newfilename + i + "\\\\"
    newfilename = newfilename + sha256
    """如果使用sha256命名的文件已存在，说明两个文件的sha256值相同，即文件重复，删除文件；否则重命名"""
    if sha256.lower() == oldname.lower():
        pass
    else:
        if os.path.exists(newfilename):
            print(f"[-]--->文件已存在：{filepath}")
            os.remove(filepath)
            return 1
        else:
            os.rename(filepath, newfilename)
            filepath = newfilename

    """输出压缩文件的位置，即将压缩文件输出到哪个文件夹中，压缩文件以文件的 SHA256.rar 命名"""
    if os.path.exists(outputpath):
        pass
    else:
        os.makedirs(outputpath)

    outputpath = f"{outputpath}\\{sha256}"

    """如果压缩文件已存在，则说明录入重复，删除源文件即可"""
    if os.path.exists(outputpath + ".rar"):
        print(f"[-]--->压缩文件已存在:\t{filepath}")
        os.remove(filepath)
        return 1

    """DOS命令"""
    cmdzip = f"{WinRARpath} a -ep -p{password}  {outputpath} \"{filepath}\""
    try:
        # DOS调用WinRAR加密压缩文件
        os.popen(cmdzip)
        if filesize > 5000000:
            time.sleep(5)
        else:
            time.sleep(3)
        os.remove(filepath)
        print(f"[+]==>源文件压缩成功：\t{filepath}")
    except Exception as err:
        print(err)


"""记录日志"""


def writeLog(hashes, apt):
    for k, v in hashes.items():
        filepath = k
        oldname = filepath.split('\\')[-1]
        md5 = v['MD5']
        sha1 = v['SHA1']
        sha256 = v['SHA256']
    with open("d:\\sample.txt", "a", encoding="utf-8")as f:
        hashs = oldname + "#" + md5 + "#" + sha1 + "#" + sha256 + "#" + apt + "\n"
        f.writelines(hashs)


"""调用DOS解压缩"""


def unzipFile(filepath, WinRARpath, password, flag):
    outputpath = "D:\\unzipfile"
    cmdunzip = f"{WinRARpath} e -p{password}  {filepath} {outputpath}"
    try:
        # DOS调用WinRAR加密压缩文件
        os.popen(cmdunzip)
        print(f"[+]==>源文件解压成功：\t{filepath}")
        if flag == 0:
            # 删除原有文件
            # os.remove(filepath)
            print(f"[+]==>源文件删除成功：\t{filepath}")
        elif flag == 1:
            pass
    except Exception as err:
        print(err)


def work(sample, WinRARpath, password, outputpath, apt):
    print(sample)
    """检测APT属性"""
    aptname = apt
    """提取hash"""
    hashes = getHash(sample)
    """获取样本的type类型"""
    """检测数据库中是否已经存在"""
    zipFile(hashes, WinRARpath=WinRARpath, password=password, outputpath=outputpath)
    """入库"""
    """记录日志"""
    writeLog(hashes, aptname)


if __name__ == '__main__':
    WinRARpath = r"D:\WinRAR\Rar.exe"
    password = "￥bA1Ze$H2O.C0m/"
    outputpath = r"E:\zipsample18"
    path = r"E:\样本_备份完整\20200413\三天\181010"

    apt = ''
    samples = getfielpath(path)
    print(f"共计样本：{len(samples)}")

    processnum = 100
    while samples:
        process_list = []
        for i in range(processnum):
            if len(samples) != 0:
                # 创建进程
                process_list.append(Process(target=work, args=(samples.pop(), WinRARpath, password, outputpath, apt,)))
        for i in process_list:
            # 开启进程
            i.start()
        for i in process_list:
            # 进程同步
            i.join()