pefile optimization (get_memory_mapped_image)

test_reloc_pefile.py

import pefile

"""
$ python test_reloc_pefile.py
Test with ntdll.dll - with bytearray and custom setter
        Parse PE : 0.054 seconds
        574 relocs to handle !
        Relocate it : 0.048 seconds
Test with ntdll.dll - with classic buffer
        Parse PE : 0.053 seconds
        574 relocs to handle !
        Relocate it : 0.909 seconds
Test with mshtml.dll - with bytearray and custom setter
        Parse PE : 11.524 seconds
        272132 relocs to handle !
        Relocate it : 89.830 seconds
Test with mshtml.dll - with classic buffer
        Parse PE : 12.748 seconds
        272132 relocs to handle !
^C [...]
"""

MSHTML_DLL = "/mnt/c/Windows/System32/mshtml.dll"
NTDLL_DLL  = "/mnt/c/Windows/System32/ntdll.dll"

import time
# Stolen from stackoverflow
class benchmark(object):
    def __init__(self,name):
        self.name = name
    def __enter__(self):
        self.start = time.time()
    def __exit__(self,ty,val,tb):
        end = time.time()
        print("%s : %0.3f seconds" % (self.name, end-self.start))
        return False


class CustomPE(pefile.PE):
    # Added support to bytearray to enhance perf (when applying relocations)
    def get_bytes_from_data(self, offset, data):
        """."""
        if offset > len(data):
            return b''
        d = data[offset:]
        if isinstance(d, bytearray):
            return str(d)
        return d
    def set_bytes_at_offset(self, offset, data):
        """Overwrite the bytes at the given file offset with the given string.
        Return True if successful, False otherwise. It can fail if the
        offset is outside the file's boundaries.
        """
        if offset >= 0 and offset < len(self.__data__):
            self.__data__[offset:offset + len(data)] = data
        else:
            return False
        return True
        
def test(dll, custom):
    data = bytearray(open(dll, 'rb').read())
    dllname = dll.split('/')[-1]
    
    print "Test with {} - {}".format(dllname, "with bytearray and custom setter" if custom else "with classic buffer")
    
    with benchmark("\tParse PE"):
        if custom:
            pe = CustomPE(data=data)
        else:
            pe = pefile.PE(data=data)
    
    print "\t{} relocs to handle !".format(sum([len(i.entries) for i in pe.DIRECTORY_ENTRY_BASERELOC]))
    
    with benchmark("\tRelocate it"):
        x = pe.get_memory_mapped_image(ImageBase=0xdead0000)
    
    return x


test(NTDLL_DLL, 1)
test(NTDLL_DLL, 0)

test(MSHTML_DLL, 1)
# Never end ...
test(MSHTML_DLL, 0)