Realign the sections of a PE file with invalid FileAlignment.

realign_pe_sections.py

#!/usr/bin/env python2
'''
some documentation
author: Willi Ballenthin
email: [email protected]
website: https://gist.github.com/williballenthin/d43cbc98fa127211c9099f46d2e73d2c
'''
import sys
import logging
from collections import namedtuple

import pefile
import hexdump
import argparse


logger = logging.getLogger(__name__)


DEFAULT_FILE_ALIGNMENT = 0x200
DEFAULT_SECTION_ALIGNMENT = 0x1000


RealignedSection = namedtuple('RealignedSection',
                              ['name', 'src_offset', 'src_size', 'src_buf', 'dst_offset', 'dst_size', 'dst_padding'])


def align(offset, alignment):
    """
    Return the offset aligned to the nearest greater given alignment

    Args:
      offset (int): 
      alignment (int): 
    
    Returns:
      int:
    """
    if offset % alignment == 0:
        return offset
    return offset + (alignment - (offset % alignment))


def realign_pe(buf):
    '''
    Args:
      buf (bytes): a buffer containing a PE file
    
    Returns:
      bytes: a buffer containing a PE file
    '''
    pe = pefile.PE(data=buf)

    SIZEOF_FILE_HEADER = 0x18 
    header_size = pe.DOS_HEADER.e_lfanew + SIZEOF_FILE_HEADER + pe.FILE_HEADER.SizeOfOptionalHeader + (pe.FILE_HEADER.NumberOfSections * 0x28)
    aligned_header_size = align(header_size, 0x200)

    # fixup FileAlignment field
    logger.debug('adjusting file alignment: %s -> %s', hex(pe.OPTIONAL_HEADER.FileAlignment), hex(DEFAULT_FILE_ALIGNMENT))
    pe.OPTIONAL_HEADER.FileAlignment = DEFAULT_FILE_ALIGNMENT

    dst_offset = aligned_header_size  # the offset at which the current section should begin
    dst_secs = []  # list of RealignedSection instances
    for section in sorted(pe.sections, key=lambda s: s.PointerToRawData):
        dst_size = align(section.SizeOfRawData, 0x200)
        padding = '\x00' * (dst_size - section.SizeOfRawData)

        # collect pointers to the section data
        sec = RealignedSection(
            section.Name,
            section.PointerToRawData,
            section.SizeOfRawData,
            buf[section.PointerToRawData:section.PointerToRawData + section.SizeOfRawData],
            dst_offset,
            dst_size,
            padding)

        logger.debug('  resizing %s\toffset: %s\traw size: %s  \t--> offset: %s\traw size: %s',
                     section.Name,
                     hex(section.PointerToRawData),
                     hex(section.SizeOfRawData),
                     hex(dst_offset),
                     hex(dst_size))

        dst_secs.append(sec)

        # fixup the section pointers
        section.PointerToRawData = dst_offset
        section.SizeOfRawData = dst_size
        dst_offset += dst_size


    # get the modified data
    mod_buf = pe.write()
    assert pefile.PE(data=mod_buf).OPTIONAL_HEADER.FileAlignment == DEFAULT_FILE_ALIGNMENT

    ret = []
    ret.append(str(mod_buf[:header_size]))
    ret.append('\x00' * (aligned_header_size - header_size))

    for sec in dst_secs:
        ret.append(sec.src_buf)
        ret.append(sec.dst_padding)

    return ''.join(ret)


def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]

    parser = argparse.ArgumentParser(description="Realign the sections of a PE file with invalid FileAlignment.")
    parser.add_argument("input", type=str,
                        help="Path to input file")
    parser.add_argument("output", type=str,
                        help="Path to output file")
    parser.add_argument("-v", "--verbose", action="store_true",
                        help="Enable debug logging")
    parser.add_argument("-q", "--quiet", action="store_true",
                        help="Disable all output but errors")
    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    elif args.quiet:
        logging.basicConfig(level=logging.ERROR)
    else:
        logging.basicConfig(level=logging.INFO)


    with open(args.input, 'rb') as f:
        with open(args.output, 'wb') as g:
            g.write(realign_pe(f.read()))


if __name__ == "__main__":
    sys.exit(main())

example output:

DEBUG:__main__:adjusting file alignment: 0x10 -> 0x200
DEBUG:__main__:  resizing .text         offset: 0x310   raw size: 0x1f70        --> offset: 0x400    raw size: 0x2000
DEBUG:__main__:  resizing .rdata        offset: 0x2280  raw size: 0x810         --> offset: 0x2400   raw size: 0xa00
DEBUG:__main__:  resizing .data         offset: 0x2a90  raw size: 0x30          --> offset: 0x2e00   raw size: 0x200
DEBUG:__main__:  resizing .pdata        offset: 0x2ac0  raw size: 0x290         --> offset: 0x3000   raw size: 0x400
DEBUG:__main__:  resizing .CRT          offset: 0x2d50  raw size: 0x10          --> offset: 0x3400   raw size: 0x200
DEBUG:__main__:  resizing .rsrc         offset: 0x2d60  raw size: 0x360         --> offset: 0x3600   raw size: 0x400
DEBUG:__main__:  resizing .reloc        offset: 0x30c0  raw size: 0x50          --> offset: 0x3a00   raw size: 0x200