#!/usr/bin/env python
'''
Matthew Wollenweber
August 6 2009
mjw@cyberwart.com

I want to debug modules that are dynamically loaded at runtime and then unloaded. It's a pain to set and
reset breakpoints. This script watches for the module to be loaded. When the module is loaded breakpoints are set. 

-r implies the use of a relative base for the target module
'''
DESC = "Set breakpoints when a specific library is loaded"
USAGE = "!stalkmod <-r> modname <breakpoint 0> ... <breakpoint n>"


import immlib
from immlib import LogBpHook
#import time
import struct
import unicodedata
import getopt



def main(args): 
    imm = immlib.Debugger()
    imm.Log("starting modstalker")
    breakpoints = []
    module_name = ""
    
    
    try:
        opts, extra = getopt.getopt(args, "r")
        
        for flag, val in opts:
            if flag == "-r":
                imm.addKnowledge("getbase", 1)
        
        module_name = extra[0]
        for x in range(1, len(extra)):
            breakpoints.append(long(extra[x], 16))
    except:
        usage()
        imm.Log("ERROR in modstalker")
        return "ERROR. Exiting module"

    allmodules=imm.getAllModules()
    for x in allmodules:
        if x.getName().find(module_name) >= 0:
            for b in breakpoints:
                imm.setBreakpoint(b)        
            imm.Log("Module Already loaded. Breakpoints set")
            return "Module already loaded. Breakpoints set"
    
    imm.addKnowledge("breakpoints", breakpoints)
    imm.addKnowledge("modname", module_name)
                     
    load_hook = MyLoadHook()

    load_addr = imm.getAddress("kernel32.LoadLibraryW")
    load_hook.add("bp on loadlibraryW", load_addr, 0, 0, 0)
    load_addr = imm.getAddress("kernel32.LoadLibraryA")
    load_hook.add("bp on loadlibraryW", load_addr, 0, 0, 0)
    imm.Log("Placed hook on loadlibraryA and LoadLibraryW")
    
    return "Hooks loaded"

    
class MyLoadHook(LogBpHook):
    def __init__(self):
        LogBpHook.__init__(self)
        self.imm = immlib.Debugger()
        self.found = 0
             
    def run(self,regs):       
        imm = self.imm
        ucode = 0
        getbase = 0
        baseoffset = long(0)
        
        modname = imm.getKnowledge("modname")
        getbase = imm.getKnowledge("getbase")
        
        regs = imm.getRegs()
        ptr = imm.readMemory(regs['ESP']+0x4, 0x4)
        ptr = int(struct.unpack("L", ptr)[0])
        filename = imm.readString(ptr)
        
        
        #fuck unicode
        if len(filename) < 3:
            ucode = 1
            tmp = imm.readWString(ptr)
            filename = ""
            x = 0
            while x <  len(tmp):
                filename += str(tmp[x])
                x+=2
                
        filename = filename.lower()
        
        #imm.Log("I'm in: " + filename)
        #imm.Log("I'm stalking: " + modname)
        
        #should be imm.getFunction(addr).getFunctionEnd
        if filename.find(modname) >= 0:
            imm.Log("MATCHED. Setting breakpoints")
            if ucode == 0:
                imm.setBreakpoint(0x7c801DAB) #End of LoadLibraryA
            else:
                imm.setBreakpoint(0x7c80AEEB) #End of LoadLibraryW
              
            imm.Run() #let it run until the end of loadlibrary so that our memory is mapped
            #time.sleep(2)
            if getbase == 1:   
                allmodules=imm.getAllModules()
                for x in allmodules:
                    if x.getName().find(modname) >= 0:
                        baseoffset = x.getBase()
                        break
                
            bps = imm.getKnowledge("breakpoints")
            for x in bps:
                x = x + baseoffset
                imm.Log("set BP: %08x" % (x))
                imm.setBreakpoint(x)
            
            imm.Log("breakpoints set")
            self.found = 1
            self.UnHook()
                 
        return 
    
          
def usage():
    print USAGE


if __name__=="__main__":
    print "This module is for use within Immunity Debugger only"