pykd windbg driver trace script

driver-trace.py

# 'sxe ld driver.sys' to breakpoint when driver loads
# probably want to manually breakpoint DriverEntry then trace...
# or find a better solution than this ghetto script

# modified to use pykd api more and increase the speed significantly

import pykd

PATH = "C:\\Users\\droogie\\Desktop\\trace.txt"
MODULE_NAME = "driver.sys"
PREAMBLE = MODULE_NAME + "+"

#def get_module_base(module_name):
#    res = pykd.dbgCommand('dd {} L?1'.format(module_name))
#    return int(res.split()[0].replace("`", ""), 0x10)
#
#def get_module_range(module_name):
#    res = pykd.dbgCommand('lm m {}'.format(module_name))
#    return int(res.split("\n")[2].split()[1].replace("`", ""), 0x10)

def getCurrentAddr():
    #return int(pykd.dbgCommand('r rip')[4:], 0x10)
    return pykd.reg('rip')

def getCurrentOffset():
    return getCurrentAddr() - module_base

def logTrace():
    f.write(PREAMBLE + hex(getCurrentOffset()) + "\n")

def withinRange(address):
    if address >= module_base and address <= module_range:
        return True
    return False

def isModuleLoaded():
    #being lazy
    if len(pykd.dbgCommand('lm m ' + MODULE_NAME.strip('.sys')).split('\n')) > 3:
        return True
    return False


#module_base = get_module_base(MODULE_NAME)
#module_range = get_module_range(MODULE_NAME[:-4])

myModule = pykd.module(MODULE_NAME[:-4])
module_base = myModule.begin()
module_range = myModule.end()

f = open(PATH, "w")

#while (isModuleLoaded()):
while (True):
    if (withinRange(getCurrentAddr())):
        logTrace()
        #pykd.dbgCommand('t')
        pykd.trace()
    else:
        logTrace()
        #pykd.dbgCommand('gu')
        pykd.stepout()

    try:
        pykd.module(MODULE_NAME[:-4])
    except:
        break

f.close()