# -*- coding: utf-8 -*-
import os
import pickle
import random
import shutil
import time

from vtlite import md5sum
from vtlite import vtAPI

gPath = r'/home/ldy/download/20131011003ndb'             #sourceFile path
gClassPath = r'/home/ldy/download/newDataBase'            #the folder path  you want to cluster flle relied on the malvare familly
gSaveFilePath = r'/home/ldy/download/databaseinf2.pkl'     #the file path to save responses about malware samples from virustotal
gVirusTotalPublicApiKey = 'b6cf83f13fc9ee58cf429389fdf6502dc793f130ce93935c0a925b0bd876eb15' #your virustotal Public API KEY
normal_speciels = () #常见病毒类型
class baseInfor:  #统计病毒库的种类信息
    def __init__(self,groot):
        f = open(groot,'rb')
        #字典格式为{‘filename’：['malware_kind',prob_mal_kind,detail]}
        self._gdict = pickle.load(f)
        f.close()
        self._malware_species = self.__GetMalwareSpecies__()
        self._malwareInfor = self.__CountPerMalwareNum__()

    def __GetMalwareSpecies__(self):
        tlist = []
        for k,v in self._gdict.items():
            if v[0]:
                normalFlag = True
                for nitm in normal_speciels:
                    if nitm in v[0]:
                        if nitm not in tlist:
                             tlist.append(nitm)
                        normalFlag = False
                        break
                if normalFlag:
                    if v[0] in tlist:
                        continue
                    else:
                        tlist.append(v[0])
        return tlist

    def __CountPerMalwareNum__(self): #统计每一种病毒的样本数
        temp_list ={}
        for k,v in self._gdict.items():
            if v[0]:
                norm_flag = True
                for nitm in normal_speciels:
                    if nitm in v[0]:
                        temp = nitm
                        norm_flag = False
                        break
                if norm_flag:
                     temp = v[0]
            else:
                temp = v[0]
            for i in range(len(self._malware_species)):
                if self._malware_species[i] == temp:
                    if temp in temp_list:
                        temp_list[temp] += 1
                    else:
                        temp_list[temp] = 1
                    break
        return temp_list

    def GetMalwareSpecies(self):  #返回病毒类型，不包括非病毒文件
        result= []
        for k,v in self._malwareInfor.items():
            if k == 'None' or k == 'Unknow' or k == 'Undetected':
                continue
            else:
                result.append(k)
        return result

    def GetDefMalwareNum(self,mal): #打印指定类型的样本数,包括病毒类型和非病毒类型('Unkown','None','Undetected')
        vg = -1
        for k,v in self._malwareInfor.items():
            if k == mal:
                vg = v
                break
        return vg

    def GetSpeciesByFileName(self,fileName):
        value = self._gdict[fileName]
        return value[0]

    def PrintVirustotalInforByFileName(self,fileName):#返回指定文件的病毒类型
        print('the sample has a detect rate: ', self.gdict[fileName][1],' and it is Known by Symantac as ',self.gdict[fileName][0])
        print('the detail information collected by VirusTotal is listed by follow')
        reports = self.gdict[fileName][2]
        if len(reports) <= 1:
            print('the file is undeccted')
        else:
            for x in reports:
                print('\t', x, '\t' if len(x) < 7 else '', '\t' if len(x) < 14 else '', '\t',
                      reports[x]['detected'],
                      '\t', reports[x]['result'])
    def printmalwareInfor(self):
        print('the number of DB samples is ',len(self._gdict))
        gc = sorted(self._malwareInfor.items(),key = lambda d:d[1], reverse = True)
        for itm in gc:
            print('%-40s\t%d'%(itm[0],itm[1]))
    def getUpdateTime(self):
        mintime = 0
        maxtime = 0
        i = 0
        updateCount ={}
        for k in self._gdict:
            if len(self._gdict[k][2]) > 10:
                v = self._gdict[k][2]
                if 'Symantec' in v:
                    if(hasattr(self._gdict[k][2]['Symantec'],'update')):
                        if self._gdict[k][2]['Symantec']['update']:
                            temp = self._gdict[k][2]['Symantec']['update']
                            years =  temp[:4]
                            timet = int(temp[:8])
                            if i == 0:
                                mintime = timet
                                maxtime = timet
                            else:
                                if mintime > timet:
                                    mintime = timet
                                if maxtime < timet:
                                    maxtime = timet
                            i += 1
                            if years in updateCount:
                                updateCount[years] += 1
                            else:
                                updateCount[years] = 1

        return mintime,maxtime,updateCount

#字典格式为{‘filename’：['malware_kind_by_Symantec',detail.detect_date]}
dict_inf = {}
def report_analysis(reports,fileName):
    temp = True
    temp_v =[]
    if(reports['response_code'] == 1):#有效报告
        for x in reports['scans']:
            if x == 'Symantec':
                temp_v.append(reports['scans'][x]['result'])
                temp = False
                break
        if (temp):
            temp_v.append('Unknow')
        temp_v.append(reports['positives'] / reports['total'])  # 检测率
        temp_v.append(reports['scans']) #scans字典
    else:
        temp_v=['Undetected',0,{'Undefined':'undefined'}]
    dict_inf[fileName] = temp_v

def saveFile(groot):
    f = open(groot,'wb')
    pik = pickle.Pickler(f)
    if len(dict_inf) > 0:
        pik.dump(dict_inf)
    f.close()


def load_malware_file(PEfile_Path):   #获取所有文件的文件路径
    test_files = list()
    fnamelist = list()
    for dirpath, dirname, filenames in os.walk(PEfile_Path):
        for filename in filenames:
            rfpath = os.path.join(dirpath, filename)
            test_files.append(rfpath)
            fnamelist.append(filename)
    return test_files,fnamelist

def SampleClassfyByVirusTotal(groot):
    global gVirusTotalPublicApiKey

    fpath,fnames = load_malware_file(groot)
    api = vtAPI(gVirusTotalPublicApiKey)
    random_sampling = random.sample([i for i in range(len(fpath))],len(fpath))

    for order in random_sampling:
        if fnames[order] in dict_inf:
            continue
        md5 = md5sum(fpath[order])
        try:
            result = api.getReport(md5)
        except:
            saveFile(gSaveFilePath)
            return False
        report_analysis(result, fnames[order])
        print(order) #观察中间过程
        time.sleep(16)

    saveFile(gSaveFilePath)
    return  True

def MoveMalwareFile():
    # 将同家族的病毒文件移动到相同的文件中
    binf = baseInfor(gSaveFilePath)

    os.makedirs(gClassPath)
    fpath, fnames = load_malware_file(gPath)
    i = 0
    for fp in fpath:
        type = binf.GetSpeciesByFileName(fnames[i])
        tempPath = gClassPath + '\\' +  type
        if(not os.path.exists(tempPath)):
            os.makedirs(tempPath)
        try:
            shutil.move(fp,tempPath)
        except Exception:
            print('moving file meet error')
            break
        i += 1

def LoadPreDetectInf(): #载入已有的病毒信息数据库
    global dict_inf
    if os.path.exists(gSaveFilePath):
        f = open(gSaveFilePath,'rb')
        #字典格式为{‘filename’：['malware_kind',prob_mal_kind,detail]}
        dict_inf  = pickle.load(f)
        f.close()

def main():
    #LoadPreDetectInf()
    #while(1):                                  #等待分类完所有样本后，终止程序
     #   if SampleClassfyByVirusTotal(gPath):    #通过virustotal鉴别文件
       #     break
      #  else:
       #     time.sleep(5)#发生错误时，停止工作一段时间，再重新发起请求
        #    LoadPreDetectInf()
    binf = baseInfor(gSaveFilePath)
    mspecies = binf.GetMalwareSpecies()
    print('there are %d kinds of Malware in the datebase'%len(mspecies))
    binf.printmalwareInfor()
    mtime ,mxtime,tcount = binf.getUpdateTime()
    print('the time range from %d to %d'%(mtime,mxtime))

    for its in tcount:
        print('year:%-10s%d'%(its,tcount[its]))

    #for species in mspecies:
     #   print(species,':',binf.GetDefMalwareNum(species))

    #MoveMalwareFile()         #

if __name__ == '__main__':
    main()

