Python 扫描PE文件头信息

'''
Python 扫描PE文件头信息 by 郑瑞国
getPEinfo.py
'''
import os
import time
import string
import hashlib
import pefile
import datetime
import threading
 
def gethash(file):
    m = hashlib.md5()
    s = hashlib.sha1()
    s256 = hashlib.sha256()
    with open(file,'rb') as f:
        for line in f:
            m.update(line)
            s.update(line)
            s256.update(line)
    md5code = m.hexdigest()
    sha1code = s.hexdigest()
    sha256code = s256.hexdigest()
    return (md5code,sha1code,sha256code)
 
def getdisklist():
    disklist = []
    d = string.ascii_uppercase
    #print(d)
    for w in d:        
        disk = w+':'        
        if os.path.isdir(disk):            
            disklist.append(disk)
    return disklist

def getPEinfo(myfile):
    try:
        pe = pefile.PE(myfile)
        warning = pe.get_warnings()
        mymd5code,mysha1code,mysha256code = gethash(myfile)
        if warning:
            print()
            with open ('d:/md5.txt','a') as f:
                f.write(myfile+'\n'+mymd5code+'\n'+str(warning)+'\n')
            print(myfile)
            print('MD5: ',mymd5code)
            print('SHA-1:',mysha1code)
            print('SHA-256:',mysha256code)
            print('File Name:',os.path.basename(myfile))
            print('File Size:',os.path.getsize(myfile),'byte')
            print('Optional Header:',hex(pe.OPTIONAL_HEADER.ImageBase))
            print('EntryPoint:',pe.OPTIONAL_HEADER.AddressOfEntryPoint)
            print('Compile Time:',datetime.datetime.fromtimestamp(pe.FILE_HEADER.TimeDateStamp))
            print('Subsystem:',pefile.SUBSYSTEM_TYPE[pe.OPTIONAL_HEADER.Subsystem])
            print('DLL:',pe.FILE_HEADER.IMAGE_FILE_DLL)
            print('Sections:',pe.FILE_HEADER.NumberOfSections)
            print('warning:')
            for w in warning:
                print(w)
            for importdll in pe.DIRECTORY_ENTRY_IMPORT:
                print(importdll.dll.decode())
    except:
        pass
 
def scan(disklist):
    for disk in disklist:
        os.chdir(disk+'/')
        tree = os.walk('/')
        for dir in tree:
            for file in dir[2]:
                exname = os.path.splitext(file) 
                if '.exe'in exname[1] or '.dll'in exname[1]:
                    myfile = disk+dir[0]+'/'+file
                    getPEinfo(myfile)
                    
if __name__=='__main__':
    disklist = getdisklist()
    scan(disklist)