#!c:\Python25\python.exe
#(c) copyright thomas.pollet@gmail.com

from hooking import *
from pydbg import *
from pydbg.defines import *
from pydbg.pydbg_core import *
import getopt, sys, os, time, subprocess

################################################################################
# Tdbg : pydbg + some stuff needed for this application
################################################################################

class Tdbg(pydbg,pydbg_core):

    def __init__(self):
        
        pydbg.__init__(self)
        self.load_dll_actions = {None:[]}


    def event_handler_load_dll(self):
        '''
            This LOAD_DLL_DEBUG_EVENT handler calls all functions in self.load_dll_actions on self
        '''
        
        pydbg.event_handler_load_dll(self)
        
        last_dll = self.system_dlls[-1].name
        actions  = [] + self.load_dll_actions[None] 
        
        if last_dll in self.load_dll_actions :
            for action in self.load_dll_actions[last_dll]:
                actions.append(action)
            del self.load_dll_actions[last_dll]

        for act in actions :
            act(self)
            
        return DBG_CONTINUE


    def on_load_dll(self, proc, dll_name = None):
        '''
            Tell self to execute proc when dll with dll.name == dll_name is loaded
        '''

        if dll_name in self.load_dll_actions:
            self.load_dll_actions[dll_name].append(proc)
        else:
            self.load_dll_actions[dll_name] = [proc]


    def stack_unwind(self, do_print = False ):
        '''
            Unwind and print the call stack
        '''
        
        cstack = pydbg.stack_unwind(self)
        if do_print:
            for add in cstack:
                if self.addr_to_dll(add):
                    print '0x%08x from %s,'%(add, self.addr_to_dll(add).name),
                    print 'base at 0x%08x'%self.addr_to_dll(add).base
        return cstack


    def bp_set (self, address, description = '' , restore = True, handler = None, dll_name = None):
        '''
            pydbg.bp_set extended to handle breakpoints in dlls that are not yet loaded
        '''
        
        def bp_set_dll(dbg):
            for dll in dbg.system_dlls:
                if dll.name == dll_name :
                    print 'Setting breakpoint in %s (base : 0x%08x ) at offset 0x%08x'%(dll_name,dll.base,address)
                    print dll.path , '%x'%(dll.base + address)
                    for i in dbg.disasm_around( dll.base + address ):
                        print "%x %s"%i
                    dbg.bp_set( dll.base + address, description, restore, handler )
                    break

        if dll_name:
            self.on_load_dll(bp_set_dll, dll_name)
        else:
            pydbg.bp_set(self, address, description, restore, handler )



################################################################################
# Hooking routines
################################################################################


def alloc_fail(dbg):
    print '> 0x%08x'%dbg.context.Eip, dbg.disasm(dbg.context.Eip)
    dbg.target_thread = dbg.h_thread
    dbg.hc = hook_container()
    for alloc in 'RtlAllocateHeap', 'RtlReAllocateHeap':
        add = dbg.func_resolve('ntdll',alloc) 
        dbg.hc.add(dbg, add, 0, alloc_entry, alloc_return)
    dbg.bp_del(dbg.context.Eip)
    return DBG_CONTINUE
    

def alloc_entry(dbg,args): 
    return DBG_CONTINUE


def alloc_return(dbg,args,retval):
    if dbg.h_thread == dbg.target_thread:
        print 'At', hex(dbg.context.Eip) , 
        print 'Change ret : 0x%x >> 0' % (dbg.context.Eax)
        dbg.stack_unwind(do_print=True)
        #os.system('start kill.cmd %i '%dbg.pid)
        dbg.set_register('EAX',0)
        remove_alloc_hooks(dbg)
    elif dbg.target_thread not in dbg.enumerate_threads():
        print 'Target thread no longer exists, no fail'
        remove_alloc_hooks(dbg)
    return DBG_CONTINUE


def remove_alloc_hooks(dbg):
    for alloc in 'RtlAllocateHeap', 'RtlReAllocateHeap':
        add = dbg.func_resolve('ntdll',alloc) 
        dbg.hc.remove(dbg, add)
    

def set_alloc_monitor(dbg):
    dbg.bp_del_all()
    dbg.hits = []
    print 'base', hex(dbg.system_dlls[-1].base)
    for alloc in 'RtlAllocateHeap', 'RtlReAllocateHeap':
        add = dbg.func_resolve('ntdll',alloc) 
        print alloc , '\tat 0x%08x' % add  
        dbg.bp_set(add,handler=monitor_alloc)
    print '  Hit     Called from offset (dec)   Instruction'
    return DBG_CONTINUE


def monitor_alloc(dbg):
    dbg.alloc_count += 1

    if dbg.alloc_count < dbg.skip:
        return DBG_CONTINUE
    
    for add in dbg.stack_unwind() :
        dll = dbg.addr_to_dll(add)
        if dll and dll.name == dbg.target_dll :
            base = dbg.addr_to_dll(add).base
            try:
                off = dbg.disasm_around(add)[4][0] - base
                debugdir = dbg.outdir+str(off)
                if not(off in dbg.hits):
                    print '- %06i   \t %i   \t'%(dbg.alloc_count,off),
                    print dbg.disasm_around(add)[4][1]
                    dbg.hits.append(off)
                if dbg.outdir and not os.path.exists(debugdir):
                    os.mkdir(debugdir)
            except:        
                print 'unable to disassemble address : 0x%08x'%add
    
    return DBG_CONTINUE


def print_dll(dbg):
    dll = dbg.system_dlls[-1]
    print 'Loading',dll.name,dll.path


def av_handler(dbg):
    
    sys.stderr.write('Access violation %s\n'%dbg.outdir)
    
    for i in dbg.disasm_around(dbg.context.Eip, 16):
      print i
    print dbg.dump_context()
    
    try:
        mod = dbg.addr_to_dll(dbg.context.Eip)
        print 'mod path : ', mod.path
        print 'Eip module base : 0x%08x'%mod.base
        print 'Eip offset : 0x%08x '%(dbg.context.Eip - mod.base)
    except:
        pass
    print 'call stack :'
    dbg.stack_unwind(True)
    dbg.terminate_process()
    return DBG_CONTINUE


################################################################################
# Main etc. 
################################################################################

    
def start(dbg,address):
    dbg.load(dbg.program,dbg.arglist)
    if address and dbg.target_dll:
        # Adress is specified, malloc following address will fail
        dbg.bp_set(address=address, handler=alloc_fail, dll_name=dbg.target_dll)
    dbg.run()


def process_dir(dbg,dirname):
    acalls = [ int(d) for d in os.listdir(dirname) ]
    for call in acalls:
        print 'failing call from %i'%call
        start(dbg,call)
        

def usage():
    print 'Usage: '
    print sys.argv[0] , '[ OPTIONS ] program arguments'
    print 'OPTIONS : '
    print '\n\t-a,  --address==address \n\t\t address'
    print '\n\t-h,  --help\n\t\t print help'
    print '\n\t-m,  --monitor \n\t\t monitor what dlls are being loaded'
    print '\n\t-o,  --outdir==dirname \n\t\t output directory'
    print '\n\t-t,  --target==dllname \n\t\t test target dllname'
    print '\n\t-v,  --verbose\n\t\t verbode output'
    sys.exit(2)


if __name__ == '__main__' :

    dbg      = Tdbg()
    address  = 0
    indir    = None
    dbg.skip = 0
    dbg.outdir = ''
    dbg.target_dll = None
    
    try:
        options, remainder = getopt.gnu_getopt(sys.argv[1:], 'a:t:hms:o:vi:',['address','target','help','monitor','skip','outdir','verbose'])   
        for opt, arg in options:
            if opt in ('-t', '--target'):
                dbg.target_dll = arg
            elif opt in ('-a', '--address'):
                address = int(arg)
            elif opt in ('-d', '--directory'):
                indir = arg
            elif opt in ('-m', '--monitor'):
                dbg.on_load_dll(print_dll)
            elif opt in ('-o', '--outdir'):
                dbg.outdir = arg + '/'
            elif opt in ('-s', '--skip'):
                dbg.skip = int(arg)
            elif opt in ('-i','--indir'):
                indir = arg
            elif opt in ('-h', '--help'):
                usage()
            
        dbg.program  = remainder[0]
        dbg.arglist  = remainder[1]
    
    except:
        sys.stderr = sys.stdout
        usage()

    if not(dbg.program and dbg.arglist):
        usage()
        
    print 'Loading ', dbg.program , dbg.arglist 
    
    if dbg.target_dll:
        print 'Using dll ', dbg.target_dll
        dbg.alloc_count = 0
        if not address:
            dbg.on_load_dll(set_alloc_monitor, dbg.target_dll )
            if dbg.outdir and not os.path.exists(dbg.outdir):
                os.mkdir(dbg.outdir)
        
    if dbg.skip:
        print 'Skipping %i calls'%dbg.skip
    dbg.set_callback(EXCEPTION_ACCESS_VIOLATION, av_handler)
    
    if indir:
        process_dir(dbg,indir)
    else:
        start(dbg,address)



'''

example use:

TESTFILE=test.swf
PROGRAM=firefox
DLL=NPSWF32.dll

# List allocations being called
python atest.py $PROGRAM $TESTFILE -t $DLL

$ python atest.py $PROGRAM $TESTFILE -t $DLL -o testdir
Loading  firefox test.swf
Using dll  NPSWF32.dll
base 0x6e7b0000
RtlAllocateHeap         at 0x779a2f1f
RtlReAllocateHeap       at 0x779b59e7
  Hit     Called from offset (hex, dec)   Instruction
- 001359 0x002bdb25 2874149             call [0x6ea77378]
- 001359 0x002baea5 2862757             call 0x6ea6db10
..[SNIP]..
- 001539 0x002bbfbb 2867131             call [0x6ea77190]
- 001539 0x002bdcc3 2874563             call [0x6eaf1598]


ADDRESS=001539
# Fail at malloc called from $DLL at address $ADDRESS
python atest.py $PROGRAM $TESTFILE -t $DLL -a $ADDRESS 


# List modules being loaded
python atest.py $PROGRAM $TESTFILE -m  | tee modules.log

# Process all addresses in the testing directory (for each address found with '-o testdir' a subdir is created )
cd ~/lab/pydbg/paper
TESTFILE=fcrash.html
PROGRAM=firefox
DLL=NPSWF32.dll
ls testdir | while read add
do
    echo 'USING ADDRESS : ' $add
    OUT=testing/$add/out
    if [ -e $OUT ] 
    then
        continue
    else
        python atest.py $PROGRAM $TESTFILE -t $DLL -a $add 2>&1| tee $OUT &
        PID=$!
        #sleep 11
        #taskkill /F /PID $PID
        wait 
    fi
done

'''