callRegTracer.py

#!/usr/bin/env python

"""
(c) hAsh, 2015 run via ImmunityDbg
"""
__VERSION__ = '0.3.1'
__AUTHOR__ = 'hAsh'

import immlib
import pefile
import string

ProgAuthor = __AUTHOR__
ProgName    = 'callRegTracer'
ProgVers    = __VERSION__

DESC= "Trace all the functions called via register" #description used by PyCommands GUI
# --------------------------------------------------------------------------
# custom classes
# --------------------------------------------------------------------------
class PeInfo:
    def __init__(self, imm):
        self.imm = imm
        self.module = imm.getModule(imm.getDebuggedName())
        self.start = self.module.getBaseAddress()
        self.end = self.start + self.module.getSize()
        self.pe = pefile.PE(name=self.module.getPath())
        self.ep = self.pe.OPTIONAL_HEADER.AddressOfEntryPoint + self.start

    def logInfo(self):
        self.imm.log("--------", highlight=2)
        self.imm.log(" PeInfo ", highlight=2)
        self.imm.log(self.imm.getDebuggedName())
        self.imm.log('[%08x - %08x]' % (self.start, self.end))
        self.imm.log("EP = %08x" % self.ep)
        self.imm.log("--------", highlight=2)

    def isMyModule(self, address):
        curr_name = self.imm.getDebuggedName()
        module = self.imm.findModule(address)
        if not module:
            return False
        if curr_name == module[0].lower():
            return True
        return False

    def countMyAddresses(self, addresses):
        count = 0
        for addr in addresses:
            if self.isMyModule(addr):
                count += 1
        return count
    
    def fetchRegCalls(self):
      calls_set = set()
      regs32 = ["eax", "ebx", "ecx", "edx", "esp", "ebp", "esi", "edi"]
      for reg in regs32:
        data = "call "+ reg
        asm = self.imm.assemble(data)
        results = self.imm.search(asm)
        count = self.countMyAddresses(results)
        self.imm.log("Calls via %s: total: %d in current module: %d" % (reg.upper(), len(results), count), highlight=2)
        for addr in results:
            if self.isMyModule(addr):
              calls_set.add(addr)
      return calls_set
     
    def setBpOnList(self, addr_set, comment):
      for addr in addr_set:
        self.imm.setBreakpoint(addr)
        self.imm.setComment(addr, comment)
        #imm.log("%s  %0.8x" % (data, addr))
      self.imm.log("")
 
# --------------------------------------------------------------------------
# custom functions
# --------------------------------------------------------------------------
def listModules(imm):
    imm.log("Modules:", highlight=2)
    curr_name = imm.getDebuggedName()
    modules = imm.getAllModules()
    for key in modules.keys():
        mark = ""
        if (key == curr_name):
            imm.log("%s" % (key), highlight=1)
        else:
            imm.log("%s" % (key))
    imm.log("")

def getRegValue(imm, reg_name):
    reg_name = reg_name.upper()
    regs = imm.getRegs()
    if not regs:
        return None
    reg_value = regs[reg_name]
    #imm.log("%s = 0x%08X " % (reg_name, reg_value))
    return reg_value

def is_printable(my_string):
    for c in my_string:
        if ord(c) < 0x20 or ord(c) > 0x7e:
            return False
        return True

def printCallStack(imm, calledAt):
    stacks = imm.callStack()
    if len(stacks) == 0:
        return getParams(imm, calledAt)
    for st_arg in stacks:
        arg_dump = st_arg.getProcedure()
        if "Includes" in arg_dump:
            continue
        imm.log("> %s" % arg_dump)

def getParams(imm, calledAt):
    esp = getRegValue(imm, 'esp')
    count = 0
    while (count < 4):
        dw_param = imm.readLong(esp)
        str_param = imm.readString(dw_param)
        if not str_param:
            str_param = ""
        if not is_printable(str_param):
            str_param = ""
        imm.log("%08x %s" % (dw_param, str_param), calledAt, highlight=1)
        esp += 4
        count += 1


def printFunction(imm, calledAt, address):
    function = imm.getFunction(address)
    if function:
        imm.log("%s" % (function.getName()), calledAt, highlight=2)
        printCallStack(imm, calledAt)

def banner(imm):
    imm.log("--------------------", highlight=1)
    imm.log("%s v%s By %s" % (ProgName, ProgVers, ProgAuthor), highlight=1)
    imm.log("--------------------", highlight=1)
    
# --------------------------------------------------------------------------
# ImmunityDbg API
# --------------------------------------------------------------------------
def main(args):
    imm = immlib.Debugger()
    banner(imm)
    if imm.isFinished():
        imm.log("Restarting process...")
        imm.restartProcess()

    listModules(imm)

    peInfo = PeInfo(imm)
    call_set = peInfo.fetchRegCalls()
    peInfo.setBpOnList(call_set, "Call via register")
    
    while (True):
        address = getRegValue(imm, 'eax')
        if not address:
            break
        imm.stepIn()
        printFunction(imm, imm.getCurrentAddress(), address)
        imm.run()
    #ret is the string shown at status bar
    return ".end"

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