# -*- coding: utf-8 -*-
import sys
import zipfile
import os
sys.path.append(r'\home\ldy\download\theZoo') #添加系统路径，导入THE ZOO的Imports模块
from imports import db_handler

gZooPath = r'/home/ldy/download/theZoo'

def getZooFilePath():
    db = db_handler.DBHandler()
    details = db.get_full_details()
    filePath = list()                     #病毒文件解压后的文件路径
    file_mal_familly = list()             #与文件路径对应的病毒类型
    for itms in details:
        if itms[6] == 'bin'and (itms[9] == 'win32' or itms[9] == 'win64'):
            templist = file_extract(itms[1])
            filePath.extend(templist)
            for i in range(len(templist)):
                file_mal_familly.append(itms[2])
    return filePath,file_mal_familly

def printInfAboutTheZoo():
    result_dict = {}
    db = db_handler.DBHandler()
    details = db.get_full_details()
    tcount = 0
    for itms in details:
        if itms[2] in result_dict:
            result_dict[itms[2]] += 1
        else:
            result_dict[itms[2]] = 1
        if(itms[6] == 'bin'):
            tcount += 1
    print('the zoo has %d kinds of Malware'%len(result_dict))
    for k,v in result_dict.items():
        print(k,v)
    print('there are %d .exe'%tcount)

def file_extract(filePath):
    filePath = filePath.replace('/' ,'\\')
    rfilePath = os.path.join(gZooPath,filePath)
    zfilePath = ''
    glist = list()
    passw = b''
    for dirpath, dirname, filenames in os.walk(rfilePath):
        for filename in (f for f in filenames if f.endswith('.zip')):
            zfilePath = os.path.join(dirpath,filename)
            break
        for filename in (f for f in filenames if f.endswith('.pass')): #获取解压密码
            f = open(os.path.join(dirpath,filename),'r')
            try:
                passw = f.read()[:-1]
                passw.strip()
            finally:
                f.close()
            break
    if zfilePath.__len__() > 0:
        try:
            zfile = zipfile.ZipFile(zfilePath, 'r', zipfile.zlib.DEFLATED)
        except Exception:
            print('bad zip file',filePath)
            return []
        for filename in zfile.namelist():
            if '/' in filename:
                continue
            if filename.endswith('.dll'):
                continue
            newPath = os.path.join(rfilePath,filename)
            newPath = newPath.replace('.exe', '')
            try:
                data = zfile.read(filename, pwd = 'infected'.encode())
            except Exception:
                try:
                    data = zfile.read(filename, pwd='malware'.encode())
                except Exception:
                    print('password error     ',filePath)
                    return []
            e_magic = data[:2]
            if (e_magic != b'MZ'):
                print('not a exe')
                continue
          #  if not os.path.exists(newPath):
           #     os.makedirs(newPath)
            file = open(newPath, 'wb')
            file.write(data)
            file.close()
            glist.append(newPath)
        zfile.close()
    else:
        print(filePath)
    if(len(glist) > 1):
        print('the zip file have more than 1 object',filePath)
    return glist
