VAD Tools Code

#!/usr/bin/env python

from memutil import virt2phys, PagedOutException
from struct import unpack, calcsize

VAD_FMT = "<LLLLLL"

class InvalidVAD(Exception):
    """Exception used to indicate that some portion of tree creation failed."""
class DeleteMe(Exception):
    """Exception raised by lower level to indicate that subtree cannot be created."""

class VAD:
    def __init__(self, parent, pdba, memdump, address, level=0, allow_invalid=False):
        """Create a new VAD object. Will recursively load children.
        Arguments:
            parent - VAD object that is this VAD's parent
            pdba - Page Directory Base for this process
            memdump - file pointer to memory dump
            address - (virtual) address of VAD entry
            level - recursion level (used internally, usually no reason to set this)
            allow_invalid - force the generation of the tree, even if some subtrees
                    are paged out. Otherwise an InvalidVAD exception will be raised.
        """
        #print "DEBUG: recursion level %d" % level
        self.level = level
        
        self.parent = parent
        self.pdba = pdba
        self.memdump = memdump
        self.address = address
        try:
            address_real = virt2phys(memdump, pdba, address)
            #print "DEBUG: reading VAD at %x (%x real)" % (address, address_real)
        except PagedOutException:
            if allow_invalid:
                print "WARN: subtree paged out, returning."
                raise DeleteMe()
            else:
                raise InvalidVAD("VAD tree is invalid; some subtree"
                                 " is paged out (depth: %d)" % self.level)

        memdump.seek(address_real)
        #self.pool_header = PoolHeader(memdump.read(8))
        
        vad_size = calcsize(VAD_FMT)
        (start_address, end_address, parent_addr,
         left_addr, right_addr, flags) = unpack(VAD_FMT, memdump.read(vad_size))

        # Sanity check: parent's address is correct        
        if self.parent is not None:
            try: assert self.parent.address == parent_addr
            except AssertionError:
                print "WARN: Parent addresses do not match: expected %x, got %x" % (self.parent.address,
                                                                                    parent_addr)

        self.start_address = (start_address << 12)
        # Based on the source code in Undocumented Windows Internals, it
        # appears that end_address specifies the base of the last page
        # used by the process (the calculation they have to list the full
        # range in bytes is ((end_address+1) << 12) - 1. I'm choosing to set
        # this to the address of the first unused byte, as this fits better
        # with python's range operations.
        self.end_address = ((end_address+1) << 12)
        self.flags = flags
        
        # Load left and right subtrees, if they exist
        if left_addr != 0:
            try: self.left = VAD(self, pdba, memdump, left_addr, level+1)
            except DeleteMe: self.left = None
        else: self.left = None
        if right_addr != 0:
            try: self.right = VAD(self, pdba, memdump, right_addr, level+1)
            except DeleteMe: self.right = None
        else: self.right = None
    def print_as_table(self):
        if self.parent is None:
            print "Parent   Left     Right    Start    End     "
        else: parent_address = self.parent.address

        if self.parent: parent_address = self.parent.address
        else: parent_address = 0
        if self.left: left_address = self.left.address
        else: left_address = 0
        if self.right: right_address = self.right.address
        else: right_address = 0
        
        print "%08x %08x %08x %08x %08x" % (parent_address, left_address,
                                            right_address, self.start_address,
                                            self.end_address)
        if self.left: self.left.print_as_table()
        if self.right: self.right.print_as_table()
    def print_as_tree(self):
        print " "*self.level + "%x - %x" % (self.start_address, self.end_address)
        if self.left: self.left.print_as_tree()
        if self.right: self.right.print_as_tree()
    def make_dot(self):
        # Add header
        if not self.parent:
            current_dotstring = 'digraph processtree {\n'
            current_dotstring += 'graph [rankdir = "TB"];\n'
        else:
            current_dotstring = ''

        # Add self        
        current_dotstring += ('vad_%x [label = "{ %08x - %08x }" shape = "record" color = "blue"];\n' %
                              (self.address, self.start_address, self.end_address))
        # Add subtrees
        if self.left:
            current_dotstring += 'vad_%x -> vad_%x\n' % (self.address, self.left.address)
            current_dotstring += self.left.make_dot()
        if self.right:
            current_dotstring += 'vad_%x -> vad_%x\n' % (self.address, self.right.address)
            current_dotstring += self.right.make_dot()

        # Add footer        
        if not self.parent:
            current_dotstring += '}\n'

        return current_dotstring        
    def __iter__(self):
        yield self
        if self.left:
            for node in self.left:
                yield node
        if self.right:
            for node in self.right:
                yield node
    def __repr__(self):
        return ("<VAD node %08x-%08x, flags %08x, level %d>" %
                (self.start_address,self.end_address,self.flags,
                 self.level))