patch_refs.py

#!/usr/bin/env python
# -*- coding: UTF-8 -*-

"""
Prerequisites:
    - Python 3
    - pip install pefile
Usage:
    patch_refs.py foo.dll foo_.dll 0x1020 0x4c18
"""

import sys
import pefile

def replace_references(src_filename, dst_filename, src_addr, dst_addr, rva_form=True):
    pe = pefile.PE(src_filename)

    if pe.FILE_HEADER.Machine != pefile.MACHINE_TYPE['IMAGE_FILE_MACHINE_I386']:
        raise Exception('This script currently supports 32-bit PE files only')

    if not hasattr(pe, 'DIRECTORY_ENTRY_BASERELOC'):
        raise Exception('This PE file has no relocation information')

    image_base = pe.OPTIONAL_HEADER.ImageBase

    if not rva_form:
        src_addr -= image_base
        dst_addr -= image_base
        print('Converted source VA to RVA {0:08x}, dest VA to RVA {1:08x}'.format(src_addr, dst_addr))

    for basereloc in pe.DIRECTORY_ENTRY_BASERELOC:
        for reloc in basereloc.entries:
            if reloc.type == pefile.RELOCATION_TYPE['IMAGE_REL_BASED_HIGHLOW']:
                rva = pe.get_dword_at_rva(reloc.rva) - image_base
                if rva == src_addr:
                    print('Patching reference at {0:08x}'.format(reloc.rva))
                    pe.set_dword_at_rva(reloc.rva, dst_addr + image_base)

            elif reloc.type != pefile.RELOCATION_TYPE['IMAGE_REL_BASED_ABSOLUTE']:
                print('Skipping relocation type {0}'.format(reloc.type))

    print('Writing to {0}...'.format(dst_filename))
    pe.write(dst_filename)

if __name__ == '__main__':
    if len(sys.argv) != 5:
        print('Usage: patch_refs.py <source dll> <result dll> <source rva> <destination rva>')
    else:
        src_filename = sys.argv[1]
        dst_filename = sys.argv[2]
        src_addr = int(sys.argv[3], 0)
        dst_addr = int(sys.argv[4], 0)
        replace_references(src_filename, dst_filename, src_addr, dst_addr)