Python 扫描PE文件头信息

'''
Python 扫描PE文件头信息 by 郑瑞国
getPEinfo.py
'''
import os
import string
import hashlib
import pefile
import datetime
 
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 scan(disklist):
    #print(disklist)
    for disk in disklist:
        #print(disk)
        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
                    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
                    
if __name__=='__main__':
    disklist = getdisklist()
    scan(disklist)