# -*- coding: utf-8 -*-

from PE_Classfy.feaExtractA import feature_gets
import os
import pefile
import zipfile

gFreshDBPath = r'/home/ldy/download/freshMalwaresFromDasmalwerkForTesting'


def fakeFile_checkb(fpath):
    file_object = open(fpath, 'rb')
    try:
        e_magic = file_object.read(2)
    finally:
        file_object.close()
    if (e_magic == b'MZ'):
        return True
    else:
        return False

def un_zip(file_name):
    """unzip zip file"""
    try:
        zip_file = zipfile.ZipFile(file_name, 'r', zipfile.zlib.DEFLATED)
    except Exception:
        print('bad zip file', file_name)
        return None
    f_name = file_name.replace(".zip", "")
    if os.path.isdir(f_name + "_files"):
        pass
    else:
        os.mkdir(f_name + "_files")
    t_file = f_name + "_files"
    for names in zip_file.namelist():
        zip_file.extract(names,f_name + "_files/", 'infected')
    zip_file.close()
    return t_file


def zipfile_feature_get(rfpath):
    t_file_path = un_zip(rfpath)
    if t_file_path:
        return file_extract(t_file_path)
    else:
        return []


def file_extract(filePath):
    feature = []
    gpath = []
    for dirpath, dirname, filenames in os.walk(filePath):
        for filename in filenames:
            rfpath = os.path.join(dirpath, filename)
            if filename.endswith('.zip'):
                feature.extend(zipfile_feature_get(rfpath))
                os.remove(rfpath)
            else:
                if fakeFile_checkb(rfpath):
                    try:
                        pfile = pefile.PE(rfpath)
                        feature.append(feature_gets(pfile,rfpath))
                        #gpath.append(rfpath)
                    except Exception:
                         continue
    #MD5_collection(gpath)  #获取文件的MD5值
    return feature

def fresh_DB_feature_get():
    feature = file_extract(gFreshDBPath)
    lens = len(feature)
    cla = [1 for i in range(lens)]
    return feature,cla

import hashlib
def md5sum(filename):
    fh = open(filename, 'rb')
    m = hashlib.md5()
    while True:
        data = fh.read(8192)
        if not data:
            break
        m.update(data)
    return m.hexdigest()

def MD5_collection(Fpath):
    f = open('freshdbMD5.txt','w+')
    for fp in Fpath:
        f.writelines(str(md5sum(fp)) + '\n')
    f.close()


