#!/usr/bin/python
# -*- coding: utf-8 -*-
'''
    统计API组合出现的频度，（3个API）
'''
from __future__ import division
from PE_Classfy.PE_detect import load_file
from PE_Classfy.PE_detect import load_malware_file
from PE_Classfy.PE_detect import isTragen
from vt_detect.sample_classfy_byVTAPI import baseInfor


import  pefile

high_freq_api = [
        "accept",
		"AddCredentials",
		"bind",
        'CopyFile',
        'CreateFile',
		"CertDeleteCertificateFromStore",
		"CheckRemoteDebuggerPresent",
        'CheckSumMappedFile',
		"CloseHandle",
		"closesocket",
		"connect",
		"ConnectNamedPipe",
		"CreateMutex",
		"CreateProcess",
		"CreateToolhelp32Snapshot",
		"CreateFileMapping",
		"CreateRemoteThread",
		"CreateDirectory",
		"CreateService",
		"CreateThread",
		"CryptEncrypt",
        "DeleteDirectory",
        'DeleteFile',
		"DeleteCriticalSection",
		"DeviceIoControl",
		"DisconnectNamedPipe",
		"DNSQuery",
		"EnumProcesses",
		"ExitProcess",
		"ExitThread",
		"FindWindow",
		"FindResource",
		"FindFirstFile",
		"FindNextFile",
		"FltRegisterFilter",
		"FtpGetFile",
		"FtpOpenFile",
		"GetCommandLine",
		"GetComputerName",
		"GetCommandLineA",
		"GetCurrentProcess",
		"GetThreadContext",
		"GetDriveType",
		"GetFileSize",
		"GetFileAttributes",
		"GetHostByAddr",
		"GetHostByName",
		"GetHostName",
		"GetModuleHandle",
		"GetModuleFileName",
		"GetProcAddress",
		"GetStartupInfo",
		"GetSystemDirectory",
		"GetTempFileName",
		"GetTempPath",
		"GetTickCount",
		"GetUpdateRect",
		"GetUpdateRgn",
		"GetUserNameA",
		"GetUrlCacheEntryInfo",
		"GetVersionEx",
		"GetWindowsDirectory",
		"GetWindowThreadProcessId",
		"HeapAlloc",
		"HttpSendRequest",
		"HttpQueryInfo",
        '_initterm',
		"IcmpSendEcho",
		"IsBadReadPtr",
		"IsBadWritePtr",
		"IsDebuggerPresent",
		"InitializeCriticalSectionAndSpinCount",
		"InternetCloseHandle",
		"InternetConnect",
		"InternetCrackUrl",
		"InternetQueryDataAvailable",
		"InternetGetConnectedState",
		"InternetOpen",
		"InternetQueryDataAvailable",
		"InternetQueryOption",
		"InternetReadFile",
		"InternetWriteFile",
		"LdrLoadDll",
		"LoadLibrary",
		"LoadLibraryA",
		"LockResource",
		"listen",
		"lstrcmp",
        'MapFileAndCheckSum',
		"MapViewOfFile",
        'MAPISendMail'
        'MAPIReadMail'
		"MessageBox",
		"OutputDebugString",
        'OutputDebugStringA',
		"OpenFileMapping",
		"OpenMutex",
		"OpenProcess",
		"Process32First",
		"Process32Next",
		"recv",
		"ReadFile",
		"ReadProcessMemory",
		"RegCloseKey",
		"RegCreateKey",
		"RegDeleteKey",
		"RegDeleteValue",
		"RegEnumKey",
		"RegOpenKey",
		"ReleaseMutex",
		"RemoveDirectory",
		"send",
		"sendto",
		"SetFilePointer",
		"SetKeyboardState",
		"SetWindowsHook",
		"ShellExecute",
		"Sleep",
		"socket",
		"StartService",
		"TerminateProcess",
		"UnhandledExceptionFilter",
		"URLDownload",
		"VirtualAlloc",
		"VirtualFree",
		"VirtualProtect",
		"WaitForSingleObject",
		"WriteProcessMemory",
		"WriteFile",
		"WSASend",
		"WSASocket",
		"WSAStartup",
        'WinExec',
		"ZwQueryInformation"

        'DisableThreadLibraryCalls',
        'GetAsyncKeyState',
        'RtlDecompressBuffer',
        'OpenService',
        'GetTickCount',
        'GetSystemTime'
]
lena = len(high_freq_api)

count_dict = {}

benign_nums = 0
trojan_nums = 0
mal_nums = 0

def init_count_api_dict():
    global count_dict
    for i in range(0,lena-2):
        for j in range(i+1,lena -1):
            for k in range(j+1,lena):
                key = '%d_%d_%d' % (i, j, k)
                count_dict[key] = [0,0,0]

def count_dict_func(feat,flag):
    global count_dict
    if len(feat) != lena:
        return
    for i in range(0,lena-2):
        if feat[i] == 0:
            continue
        else:
            for j in range(i+1,lena -1):
                if feat[j] == 0:
                    continue
                else:
                    for k in range(j+1,lena):
                        if feat[k] != 0:
                            key = '%d_%d_%d'%(i,j,k)
                            if flag == 0:
                                count_dict[key][0] += 1
                            elif flag == 1:
                                count_dict[key][1] += 1
                            elif flag == 2:
                                count_dict[key][2] += 1

def API_inf_get(pfiles):
    feat = [0 for i in range(lena)]
    if (hasattr(pfiles, 'DIRECTORY_ENTRY_IMPORT')):
        for importeddll in pfiles.DIRECTORY_ENTRY_IMPORT:
            for importedapi in importeddll.imports:
                if (importedapi.name is not None):
                    if not isinstance(importedapi.name, str):
                        funcName = importedapi.name.decode()
                    else:
                        funcName = importedapi.name
                    m = 0
                    for fq_api in high_freq_api:
                        gclen = len(funcName) - len(fq_api)
                        if gclen < 0 or gclen > 2:
                            m += 1
                            continue
                        if fq_api in funcName:
                            feat[m] = 1
                            break
                        m += 1
    return feat


PEfile_normal_Path = r'/home/ldy/download/benign'
PEfile_melicious_Path = r'/home/ldy/download/20131011003ndb'


def get_g_data(train_benign_num,train_mal_num):

    import random
    global benign_nums,trojan_nums,mal_nums

    gcroot_normal = load_file(PEfile_normal_Path)
    gcroot_melicious = load_malware_file(PEfile_melicious_Path)
    binf = baseInfor(r'/home/ldy/download/databaseinf2.pkl')
    print 'complete loading'

    templist = [i for i in range(len(gcroot_normal))]
    normal_sample = random.sample(templist, train_benign_num)  # 选择正常的训练样本

    templist = [i for i in range(len(gcroot_melicious))]
    malware_sample = random.sample(templist, train_mal_num)  # 选择正常的训练样本

    for i in normal_sample:
        print i
        try:
            p_files = pefile.PE(gcroot_normal[i], True)
            feat = API_inf_get(p_files)
            count_dict_func(feat, 0)
            benign_nums += 1
            p_files.close()
        except pefile.PEFormatError:
            continue

    for i in malware_sample:
        print i
        try:
            p_files = pefile.PE(gcroot_melicious[i], True)
            jret = isTragen(gcroot_melicious[i], binf)
            flag = 1
            if jret == 2:
                trojan_nums +=1
                flag = 2
            else:
                mal_nums += 1
            feat = API_inf_get(p_files)
            count_dict_func(feat, flag)
            p_files.close()
        except pefile.PEFormatError:
            continue

def sort_dict(flag = 0):
    def vg(v):        #比较三类文件API序列差异
        a = abs(v[0] - v[1])
        b = abs(v[0] - v[2])
        c = abs(v[1] - v[2])
        return max(a+b,a+c,b+c)

    def vgs(v):   #比较木马病毒特素的API序列差异
        c = abs(v[1] - v[2])
        return c

    if flag == 0:
        blist = sorted(count_dict.items(),key = lambda x:vg(x[1]),reverse = True)
    else:
        blist = sorted(count_dict.items(), key=lambda x: vgs(x[1]), reverse=True)

    i = 0
    for it in blist:
        i += 1
        if i < 301:
            func = it[0]
            fun_set = func.split('_')
            funcname = []
            for ig in fun_set:
                funcname.append(high_freq_api[int(ig)])
            print(funcname,' ',it[1])
        else:
            break

if __name__ == '__main__':
    init_count_api_dict()
    get_g_data(1500,4000)
    for (k,v) in count_dict.items():
        v[0] = v[0]/benign_nums
        v[1] = v[1] / mal_nums
        if trojan_nums != 0:
            v[2] = v[2] / trojan_nums
        if (v[0] < 0.01 and v[1] < 0.01 and v[2] < 0.01):
            del count_dict[k]
    sort_dict(1)
