#!/usr/bin/env python

from struct import unpack
from pefile import PE
from memutil import virt2phys, PagedOutException, offsets, read_range
from hexutil import hd

def read_section(memdump,pdba,sect,img_base):
    """Read the contents of a PE section from memory. This must be done
    a page at a time, as adjacent virtual memory addresses may reference
    completely different areas of physical memory.
    Arguments:
        memdump - file pointer to the memory dump image
        pdba - address of the page directory
        sect - a Section object from pefile
        img_base - the executable's base address (necessary to calculate
                section offsets)
    Returns: a string containing the contents of the section
    """

    accum = ''
    # DUBIOUS DECISION: should we read SizeOfRawData bytes or Misc_VirtualSize?
    section_start = img_base + sect.VirtualAddress
    section_end = img_base + sect.VirtualAddress + sect.SizeOfRawData
    return read_range(memdump,pdba,section_start,section_end)

def reconstruct_PE(memdump,pe,pdba):
    """Rebuild an on-disk PE file from an image in memory.
    Arguments:
        memdump - a file pointer to the memory dump
        pe - a PE object from pefile
        pdba - the address of the page directory for this process
    Returns: a string representing the executable
    """
    
    img_base = pe.OPTIONAL_HEADER.ImageBase
    rebuilt_exe = ''

    # Add the pe header:
    # Dubious decision: the PE header is in the first page at ImageBase
    img_base_real = virt2phys(memdump,pdba,img_base)
    memdump.seek(img_base_real)
    rebuilt_exe += memdump.read(4096)
    
    for sect in pe.sections:
        print "DEBUG: Starting section %s" % sect.Name
        section_data = read_section(memdump,pdba,sect,img_base)
        # If there is a gap, fill it with zeroes
        if sect.PointerToRawData > len(rebuilt_exe):
            rebuilt_exe += '\0' * (sect.PointerToRawData - len(rebuilt_exe))
        # If there is overlap, later section takes precedence, ie, truncate
        elif sect.PointerToRawData < len(rebuilt_exe):
            rebuilt_exe = rebuilt_exe[:sect.PointerToRawData]
        rebuilt_exe += section_data
    return rebuilt_exe

def unpack_le(str):
    """Silly helper function to convert 4 characters to a little-endian unsigned int"""
    return unpack("<L", str)[0]

if __name__ == "__main__":
    from optparse import OptionParser

    usage = "usage: %prog [options] <memory dump> <EPROCESS offset>"
    parser = OptionParser(usage=usage)
    parser.add_option("-e", "--executable", dest="exename",
                      help="write reconstructed executable out to FILE [default: [EPROCESS].exe]",
                      metavar='FILE')
    parser.add_option("-o", "--operating-system", dest="osname", default="XPSP2",
            help=("operating system memory dump comes from"
                  " [default: %%default, options: %s]" % ",".join(offsets.keys())))
    
    (options, args) = parser.parse_args()

    if len(args) != 2:
        import sys
        parser.print_help()
        sys.exit(1)

    print "Done parsing arguments."
    
    memdump = open(args[0], 'rb')
    eproc_offset = int(args[1], 0)
    offs = offsets[options.osname]
    if options.exename:
        output_file = open(options.exename, 'wb')
    else:
        output_file = open(str(eproc_offset) + '.exe', 'wb')

    memdump.seek(eproc_offset)
    eproc_struct = memdump.read(offs["EPROC_SIZE"])

    pdba = unpack_le(eproc_struct[offs["PDBA_OFFSET"]:offs["PDBA_OFFSET"]+4])

    peb_addr_virt = unpack_le(eproc_struct[offs["PEB_OFFSET"]:offs["PEB_OFFSET"]+4])
    peb_addr_real = virt2phys(memdump, pdba, peb_addr_virt)
    print "DEBUG: PEB found at %x (%x)" % (peb_addr_virt,peb_addr_real)
    memdump.seek(peb_addr_real + offs["PEB_IMG_BASE"])
    img_base_virt = unpack_le(memdump.read(4))
    img_base_real = virt2phys(memdump, pdba, img_base_virt)

    memdump.seek(img_base_real)
    pe_page = memdump.read(4096)

    pe = PE(data=pe_page, fast_load=True)
    output_file.write(reconstruct_PE(memdump, pe, pdba))
    output_file.close()
    memdump.close()