# -*- coding: utf-8 -*-
import json
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",
        "CryptDecrypt",
        "DeleteDirectory",
        'DeleteFile',
		"DeleteCriticalSection",
		"DeviceIoControl",
		"DisconnectNamedPipe",
		"DNSQuery",
		"EnumProcesses",
		"ExitProcess",
		"ExitThread",
		"FindWindow",
		"FindResource",
		"FindFirstFile",
		"FindNextFile",
		"FltRegisterFilter",
		"FtpGetFile",
		"FtpOpenFile",
		"GetCommandLine",
		"GetComputerName",
		"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",
		"LockResource",
		"listen",
		"lstrcmp",
        'MapFileAndCheckSum',
		"MapViewOfFile",
        'MAPISendMail'
        'MAPIReadMail'
		"MessageBox",
		"OutputDebugString",
		"OpenFileMapping",
		"OpenMutex",
		"OpenProcess",
		"Process32First",
		"Process32Next",
		"recv",
		"ReadFile",
		"ReadProcessMemory",
		"RegCloseKey",
		"RegCreateKey",
		"RegDeleteKey",
		"RegDeleteValue",
		"RegEnumKey",
		"RegOpenKey",
		"ReleaseMutex",
		"RemoveDirectory",
		"send",
		"sendto",
        "SetEnvironmentVariable",
        "SetFileAttributes",
		"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'
]
mal_string=['Scripting.FileSystemObject','Wscript.shell']

high_freq_resource = [11, 14, 16]

dllmode = [0x0040, 0x0080, 0x0100, 0x0200, 0x0400, 0x0800, 0x8000]
tchamode = [0x00000020, 0x00000040, 0x00000080, 0x02000000, 0x10000000, \
            0x20000000, 0x40000000, 0x80000000]


def feature_get(pfiles):
    feature = []
    #DOS Header0
    if(hasattr(pfiles,'DOS_HEADER')):
        feature.append(pfiles.DOS_HEADER.e_lfanew)
    else:
        feature.append(0)
    #PE HEADER
    feature.append(pfiles.FILE_HEADER.NumberOfSections)
    feature.append(pfiles.FILE_HEADER.NumberOfSymbols)

    #OPTIONAL HEADERD
    if (hasattr(pfiles, 'OPTIONAL_HEADER')):
        if(pfiles.OPTIONAL_HEADER.CheckSum == 0):
            feature.append(1)
        else:
            feature.append(0)

        if (pfiles.OPTIONAL_HEADER.AddressOfEntryPoint > pfiles.sections[-1].VirtualAddress):
            feature.append(1)
        else:
            feature.append(0)

        if (pfiles.OPTIONAL_HEADER.AddressOfEntryPoint < pfiles.sections[0].VirtualAddress or \
              pfiles.OPTIONAL_HEADER.AddressOfEntryPoint > (pfiles.sections[-1].VirtualAddress +pfiles.sections[-1].SizeOfRawData)):
            feature.append(1)
        else:
            feature.append(0)

        for mode in dllmode:
            feature.append(mode & pfiles.OPTIONAL_HEADER.DllCharacteristics)

        feature.append(pfiles.OPTIONAL_HEADER.SizeOfCode)

        feature.append(pfiles.OPTIONAL_HEADER.SizeOfHeapCommit)
        feature.append( pfiles.OPTIONAL_HEADER.SizeOfHeapReserve)

        feature.append( pfiles.OPTIONAL_HEADER.SizeOfStackCommit)
        feature.append(pfiles.OPTIONAL_HEADER.SizeOfStackReserve)

        feature.append(pfiles.OPTIONAL_HEADER.SizeOfInitializedData)
        feature.append(pfiles.OPTIONAL_HEADER.SizeOfUninitializedData)
        #feature.append(pfiles.OPTIONAL_HEADER.MajorImageVersion)
        #feature.append(pfiles.OPTIONAL_HEADER.MinorImageVersion)

        #feature.append(((pfiles.OPTIONAL_HEADER.MajorImageVersion * 100)+pfiles.OPTIONAL_HEADER.MinorImageVersion）*1000)

        #feature.append(pfiles.OPTIONAL_HEADER.MajorOperatingSystemVersion)
        #feature.append(pfiles.OPTIONAL_HEADER.MinorOperatingSystemVersion)

        #feature.append(pfiles.OPTIONAL_HEADER.MajorLinkerVersion)
        #feature.append(pfiles.OPTIONAL_HEADER.MinorLinkerVersion)

        if pfiles.OPTIONAL_HEADER.NumberOfRvaAndSizes > 0x10:
            feature.append(1)
        else:
            feature.append(0)

        #data directory
        tblens = len(feature)
        m = 0
        feature.extend([0 for i in range(16)])
        for dataDirect in pfiles.OPTIONAL_HEADER.DATA_DIRECTORY:
            feature[tblens + m] = dataDirect.Size
            m += 1
        #feature.append(pfiles.OPTIONAL_HEADER.DATA_DIRECTORY[1].VirtualAddress) #IATRVA

    else:
        feature.extend([0 for i in range(3 + len(dllmode) + 8 + 16)]) #去出了六个Version特征
        #feature.extend([0 for i in range( 3 +len(dllmode)+ 14 + 16 )])

    # resource
    tlens = len(feature)
    feature.extend([0 for i in range(len(high_freq_resource))])
    resid = []
    if (hasattr(pfiles, 'DIRECTORY_ENTRY_RESOURCE')):
        for rcsc in pfiles.DIRECTORY_ENTRY_RESOURCE.entries:
            temp_res_id = rcsc.id
            if (temp_res_id is None):
                temp_res_id = rcsc.struct.Id
            m = 0
            for gres_id in high_freq_resource:
                if(gres_id == temp_res_id):
                    feature[tlens + m] = 1
                m += 1
            resid.append(temp_res_id)
    NumsOfSource_mark = [2,3,6,9,11,21,23,24]
    tlens = len(feature)
    feature.extend([0 for i in range(len(NumsOfSource_mark))])
    m = 0
    for i in NumsOfSource_mark:  # 统计关键资源个数
        if (i + 1) in resid:
            for rcsc in pfiles.DIRECTORY_ENTRY_RESOURCE.entries:
                temp_res_id = rcsc.id
                if (temp_res_id is None):
                    temp_res_id = rcsc.struct.Id
                if ((i + 1) == temp_res_id):
                    if (hasattr(rcsc, 'directory')):
                        feature[tlens + m]=len(rcsc.directory.entries)
                    else:
                        feature[tlens + m] = 0
                    break
        else:
            pass
        m += 1

    # dll & api MODULE
    tlens = len(feature)
    feature.extend([0 for i in range(len(high_freq_api) + 5)])
    api_len = 0
    Excute = False
    if (hasattr(pfiles, 'DIRECTORY_ENTRY_IMPORT')):
        for importeddll in pfiles.DIRECTORY_ENTRY_IMPORT:
            if (b'mscoree.dll' == importeddll.dll):
                feature[tlens+1] = 1
            elif (b'wsock32.dll' == importeddll.dll):
                 feature[tlens+2] = 1
            elif (b'imm32.dll' == importeddll.dll):
                    feature[tlens+3] = 1
            for importedapi in importeddll.imports:
                api_len += 1
                if (importedapi.name is not None):
                    m = 0
                    for fq_api in high_freq_api:
                        funcName = importedapi.name.decode()
                        if 'ShellExecute' in funcName or 'WinExec' in funcName or 'CreateProces' in funcName:
                            Excute = True
                        if fq_api == funcName:
                            feature[tlens + 4 + m] = 1
                            break
                        m += 1
        feature[tlens] =  api_len  # 导入api个数
        if Excute:
            feature[-1] = 1

        #可疑行为
        suspicious_imports = set([b'LoadLibrary', b'GetProcAddress'])
        suspicious_imports_count = 0
        total_symbols = 0
        for imp_dll in pfiles.DIRECTORY_ENTRY_IMPORT:
            for symbol in imp_dll.imports:
                for suspicious_symbol in suspicious_imports:
                    if symbol and symbol.name and symbol.name.startswith(suspicious_symbol):
                        suspicious_imports_count += 1
                        break
                total_symbols += 1
        if suspicious_imports_count == len(suspicious_imports) and total_symbols < 20:
            feature[tlens+ len(high_freq_api) + 4] = 1
    #导出表
    if (hasattr(pfiles, 'DIRECTORY_ENTRY_EXPORT')):
         # exportNameLEN
        feature.append(pfiles.DIRECTORY_ENTRY_EXPORT.struct.NumberOfFunctions)  # ExportSize
    else:
        feature.append(0)
    # 导入section表
    tempressult = False
    for psection in pfiles.sections:
        if(psection.SizeOfRawData == 0):
            tempressult = True
            break
    if(tempressult):
        feature.append(1)
    else:
        feature.append(0)

    if (pfiles.sections[0].Name != b'.text'):
         feature.append(1)
    else:
        feature.append(0)

    if (pfiles.sections[0].Characteristics & 0x20000000) == 0:
        feature.append(1)
    else:
        feature.append(0)

    gbret2 = True
    sect_verctor = [0 for i in range(8)]
    for psection in pfiles.sections:
        if(psection.Characteristics & 0x20000000) > 0 and  (psection.Characteristics & 0x80000000) > 0 :
            if psection.Name == b'PAGE' and pfiles.is_driver():
                # Drivers can have a PAGE section with those flags set without
                # implying that it is malicious
                pass
            else:
                sect_verctor[0] += 1

        if(psection.Name == b'.text' and (psection.Characteristics & 0x20000000) == 0) or \
                (psection.Name != b'.text' and (psection.Characteristics & 0x20000000) > 0):
            sect_verctor[1] += 1

        #段名不可打印
        try:
            psection.Name.decode()
        except UnicodeDecodeError:
            sect_verctor[2] += 1

        #有可执行段
        if(psection.Characteristics & 0x20000000) > 0 and sect_verctor[3] != 0:
            sect_verctor[3] = 1

        if psection.Misc_VirtualSize > 0x10000000:
            if (sect_verctor[4] == 0):
                sect_verctor[4] = 1

    feature.extend(sect_verctor)

     #段长度大于文件总大小？？？？
    gbret2 = True
    for psection in pfiles.sections:
        if (psection.Name == b'.text' or psection.Name == b'CODE'):
            for tmode in tchamode:
                feature.append(psection.Characteristics & tmode)
            feature.append(psection.Misc_VirtualSize - psection.SizeOfRawData)
            gbret2 = False
    if(gbret2):
        feature.extend([0 for i in range(len(tchamode))])
        feature.append(0)
    if len(pfiles.sections) > 1 :
        feature.append(pfiles.sections[1].Misc_VirtualSize)
    else:
        feature.append(0)
    return feature

def feature_get_by_jiangwei(pfiles):
    feature = list()
    feature.append(pfiles.OPTIONAL_HEADER.DATA_DIRECTORY[6].Size)  # debug size
    feature.append(pfiles.OPTIONAL_HEADER.DATA_DIRECTORY[6].VirtualAddress)  # debug va

    feature.append(pfiles.OPTIONAL_HEADER.MajorImageVersion << 16 + pfiles.OPTIONAL_HEADER.MinorImageVersion)
    feature.append(
        pfiles.OPTIONAL_HEADER.MajorOperatingSystemVersion << 16 + pfiles.OPTIONAL_HEADER.MinorOperatingSystemVersion)

    feature.append(pfiles.OPTIONAL_HEADER.SizeOfStackReserve)

    feature.append(pfiles.OPTIONAL_HEADER.MajorLinkerVersion << 16 + pfiles.OPTIONAL_HEADER.MinorLinkerVersion)

    feature.append(pfiles.OPTIONAL_HEADER.DllCharacteristics)
    feature.append(pfiles.OPTIONAL_HEADER.DATA_DIRECTORY[12].VirtualAddress)  # Iat RVA
    feature.append(pfiles.OPTIONAL_HEADER.DATA_DIRECTORY[0].Size)  # ExportSize
    feature.append(pfiles.OPTIONAL_HEADER.DATA_DIRECTORY[0].VirtualAddress)  # Export Ra

    if (hasattr(pfiles, 'DIRECTORY_ENTRY_EXPORT')):
        # exportNameLEN
        feature.append(pfiles.DIRECTORY_ENTRY_EXPORT.struct.NumberOfFunctions)  # ExportSize
    else:
        feature.append(0)
    feature.append(pfiles.OPTIONAL_HEADER.DATA_DIRECTORY[2].Size)  # ResourceSize
    return feature






