#!/usr/bin/env python

from struct import unpack

# Offsets
# Add other Windows versions here as time permits
# Vista values taken from Andreas Schuster's post:
# http://computer.forensikblog.de/en/2007/01/eprocess_6_0_6000_16386.html
offsets = {
    "XPSP2": {
        "EPROC_SIZE":      0x260,  # Size of EPROCESS structure
        "PDBA_OFFSET":     0x18,   # Page Directory Base Address
        "PEB_OFFSET":      0x1B0,  # Offset to Process Environment Block
        "VAD_ROOT_OFFSET": 0x11c,  # Offset to VadRoot
        "HANDLE_TABLE_OFFSET": 0x0c4, # Offset to Handle Table
        "PEB_IMG_BASE":    0x08,    # Offset within PEB of Image Base Address
        "PEB_LDR_DATA":    0x0c    # Offset within PEB of _PEB_LDR_DATA
    },
    "2KSP4": {
        "EPROC_SIZE":      0x290,
        "PDBA_OFFSET":     0x18,
        "PEB_OFFSET":      0x1B0,
        "VAD_ROOT_OFFSET": 0x194,
        "HANDLE_TABLE_OFFSET": 0x128,
        "PEB_IMG_BASE":    0x08
    },
    "XP": {
        "EPROC_SIZE":   0x258,
        "PDBA_OFFSET":  0x18,
        "PEB_OFFSET":   0x1B0,
        "VAD_ROOT_OFFSET": 0x11c,
        "HANDLE_TABLE_OFFSET": 0x0c4,
        "PEB_IMG_BASE": 0x08
    },
    "Vista": {
        "EPROC_SIZE":   0x26C,
        "PDBA_OFFSET":  0x18,
        "PEB_OFFSET":   0x188,
        "VAD_ROOT_OFFSET": 0x238,
        "PEB_IMG_BASE": 0x08  # FIXME: UNVERIFIED (but probably correct)
    }
}

# Mask for base address in Page Directory Entries as well as Page
# Table Entries
PT_MASK = 0xFFFFF000

PD_MASK_4M = 0xFFC00000

# Page directory flags
PAGE_PRESENT = 0x00000001
PAGE_4M      = 0x00000080

# Virtual address masks
PD_ENTRY_MASK    = 0xFFC00000
PT_ENTRY_MASK    = 0x003FF000
PAGE_OFFSET_MASK_4K = 0x00000FFF
PAGE_OFFSET_MASK_4M = 0x003FFFFF

class Unsupported(Exception):
    """Exception thrown when requested operation is not suppored yet."""

class PagedOutException(Exception):
    """Empty exception class to indicate that requested page is swapped out."""

def read_range(memdump,pdba,start,end,step=4096,zero_invalid_pages=True):
    """Read the contents of a memory range. This is 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
        start - virtual address of the start of the region
        end - virtual address of the end of the region
        step - how many bytes to read at a time [default: a 4k page, or 4096 bytes]
        zero_invalid_pages - fill any invalid (paged out) areas with zeroes
    Returns: a string containing the contents of the memory range
    """

    accum = ''
    for addr in range(start, end, step):
        if end - addr < step:
            bytes_to_read = end - addr
        else:
            bytes_to_read = step
        
        try:
            real_addr = virt2phys(memdump,pdba,addr)
            #print "DEBUG: reading memory at %08x virt, %08x phys" % (addr,real_addr)
            memdump.seek(real_addr)
            accum += memdump.read(bytes_to_read)
        except PagedOutException:
            if zero_invalid_pages:
                #print "WARN: skipping memory page at 0x%08x as it is paged out." % addr
                accum += ('\0' * bytes_to_read)
            else:
                raise PagedOutException
    return accum

def virt2phys(memdump, pdba, addr):
    """Convert a virtual address to a physical one.
    A file pointer to a memory dump, the offset to a page directory,
    and a virtual address are required. The return value is a 32-bit
    unsigned integer giving the physical address of the page."""
    orig_pos = memdump.tell()
    pd_num = (addr & PD_ENTRY_MASK) >> 22
    pt_num = (addr & PT_ENTRY_MASK) >> 12
    
    
    memdump.seek(pdba + (pd_num*4))
    page_directory_entry = unpack("<L", memdump.read(4))[0]

    # Check if page table is present    
    if (page_directory_entry & PAGE_PRESENT) == 0:
        raise PagedOutException("Page table is paged out")
    # Check if it is a 4M page
    if (page_directory_entry & PAGE_4M):
        page_offset = (addr & PAGE_OFFSET_MASK_4M)
        page_base = page_directory_entry & PD_MASK_4M
        return page_base + page_offset
        #raise Unsupported("4M Pages are not yet supported")
    else:
        page_offset = (addr & PAGE_OFFSET_MASK_4K)
    
    ptba = page_directory_entry & PT_MASK
    memdump.seek(ptba + (pt_num*4))
    page_table_entry = unpack("<L", memdump.read(4))[0]

    if (page_table_entry & PAGE_PRESENT) == 0:
        raise PagedOutException("Memory page is paged out")
    page_base = page_table_entry & PT_MASK

    memdump.seek(orig_pos)
    return page_base + page_offset