enum_win_callbacks.py

'''
	Module Name: 
		enum_win_callbacks.py
	
	Abstract:
		Iterates over the nt!PspCreateProcessNotifyRoutine,
		nt!PspCreateThreadNotifyRoutine, and nt!PspLoadImageNotifyRoutine 
		callback arrays.
		
	Requirements:
		WinDbg: https://docs.microsoft.com/en-us/windows-hardware/drivers/debugger/
		PyKd: https://pykd.codeplex.com/
		
	Author:
		Nemanja (Nemi) Mulasmajic <nm@triplefault.io>
			http://triplefault.io
'''
from pykd import *
import argparse

# Discovers the size of a pointer for this system.
SIZE_OF_POINTER = (8 if (is64bitSystem()) else 4)

# The number of potential callback objects in the array.
MAXIMUM_NUMBER_OF_CALLBACKS = 0

def read_ptr(memory):
	'''
		Read a pointer of memory.
	'''
	try:
		return (ptrPtr(memory))
	except:
		print "ERROR: Could not read {} bytes from location {:#x}.".format(SIZE_OF_POINTER, memory)
		exit(-4)
		
def read_dword(memory):
	'''
		Read a DWORD of memory.
	'''
	try:
		return (ptrDWord(memory))
	except:
		print "ERROR: Could not read 4 bytes from location {:#x}.".format(memory)
		exit(-5)
		
def get_address_from_fastref(_EX_FAST_REF):
	'''
		Given a _EX_FAST_REF structure, this function will extract
		the pointer to the raw object.
	'''
	
	# kd> dt nt!_EX_FAST_REF
	#	+0x000 Object           : Ptr64 Void
	#	+0x000 RefCnt           : Pos 0, 4 Bits
	#	+0x000 Value            : Uint8B
	
	# Remove the last 4 bits of the pointer.
	return ((_EX_FAST_REF >> 4) << 4)

def enumerate_over_callbacks(array):
	'''
		Given the base of a callback array, this function will 
		enumerate over all valid callback entries.
	'''
	for i in xrange(MAXIMUM_NUMBER_OF_CALLBACKS):
		# Get the i'th entry in the array.
		entry = (array + (i * SIZE_OF_POINTER))
		entry = read_ptr(entry)
			
		# Not currently in use; skipping.
		if entry == 0:
			continue
			
		# Extract just the object pointer from the _EX_FAST_REF structure.
		callback_object = get_address_from_fastref(entry)
		
		print "{}: _EX_CALLBACK_ROUTINE_BLOCK {:#x}".format(i, callback_object)
		
		# _EX_CALLBACK_ROUTINE_BLOCK
		#	+0x000 RundownProtect   : EX_RUNDOWN_REF
		#	+0x008 Function         : PVOID
		#	+0x010 Context          : PVOID
		
		rundown_protect = read_ptr(callback_object + (SIZE_OF_POINTER * 0))
		callback_function = read_ptr(callback_object + (SIZE_OF_POINTER * 1))
		context = read_ptr(callback_object + (SIZE_OF_POINTER * 2))
		
		print "\tRundownProtect: {:#x}".format(rundown_protect)
		print "\tFunction: {:#x} ({})".format(callback_function, findSymbol(callback_function))

		type = ""
		
		if context == 0:
			type = "(Normal)"
		elif context == 2:
			type = "(Extended)"
		elif context == 6:
			type = "(Extended #2)"
		else:
			type = "(Unknown)"
			
		print "\tContext: {:#x} {}".format(context, type)

def get_address_from_symbol(mod, name):
	'''
		Attempts to locate the address of a given symbol in
		a module.
	'''
	try:
		return (mod.offset(name))
	except:
		print "ERROR: Failed to retrieve the address of {}!{}. Are symbols loaded?".format(m.name(), name)
		exit(-3)
	
parser = argparse.ArgumentParser(description='Iterates over the nt!PspCreateProcessNotifyRoutine, nt!PspCreateThreadNotifyRoutine, and nt!PspLoadImageNotifyRoutine callback arrays.')

parser.add_argument("-p", action="store_false", default=True, dest="process_callbacks", help="Exludes iteration of the process callback array (nt!PspCreateProcessNotifyRoutine).")
parser.add_argument("-t", action="store_false", default=True, dest="thread_callbacks", help="Excludes iteration of the thread callback array (nt!PspCreateThreadNotifyRoutine).")
parser.add_argument("-i", action="store_false", default=True, dest="image_callbacks", help="Excludes iteration of the image load callback array (nt!PspLoadImageNotifyRoutine).")

args = parser.parse_args()

# Must be kernel debugging to use this.
if not isKernelDebugging() and not isLocalKernelDebuggerEnabled():
	print "ERROR: This script can only be used while kernel debugging."
	exit(-1)
	
try:
	mod = module("nt")
except: 
	print "ERROR: Could not find the base address of ntoskrnl. Are symbols loaded?"
	exit(-2)
	
spacer = "=" * 100

ver = getSystemVersion()

# Vista+ can store 64 entries in the callback array.
# Older versions of Windows can only store 8.
MAXIMUM_NUMBER_OF_CALLBACKS = (64 if (ver.buildNumber >= 6000) else 8)

#####################################################################
# For process callbacks.
#####################################################################
if args.process_callbacks:
	PspCreateProcessNotifyRoutine = get_address_from_symbol(mod, "PspCreateProcessNotifyRoutine")
	PspCreateProcessNotifyRoutineCount  = read_dword(get_address_from_symbol(mod, "PspCreateProcessNotifyRoutineCount"))
	
	if ver.buildNumber >= 6001: # Vista SP1+ 
		PspCreateProcessNotifyRoutineExCount = read_dword(get_address_from_symbol(mod, "PspCreateProcessNotifyRoutineExCount"))
	else:
		PspCreateProcessNotifyRoutineExCount = 0
		
	print "\nIterating over the nt!PspCreateProcessNotifyRoutine array at {:#x}.".format(PspCreateProcessNotifyRoutine)
	print "Expecting {} nt!PspCreateProcessNotifyRoutineCount and {} nt!PspCreateProcessNotifyRoutineExCount entries.".format(PspCreateProcessNotifyRoutineCount, PspCreateProcessNotifyRoutineExCount)
	print spacer

	enumerate_over_callbacks(PspCreateProcessNotifyRoutine)

	print spacer

#####################################################################
# For thread callbacks.
#####################################################################
if args.thread_callbacks:
	PspCreateThreadNotifyRoutine = get_address_from_symbol(mod, "PspCreateThreadNotifyRoutine")
	PspCreateThreadNotifyRoutineCount = read_dword(get_address_from_symbol(mod, "PspCreateThreadNotifyRoutineCount"))
	
	if ver.buildNumber >= 10240: # Windows 10+
		PspCreateThreadNotifyRoutineNonSystemCount = read_dword(get_address_from_symbol(mod, "PspCreateThreadNotifyRoutineNonSystemCount"))
	else:
		PspCreateThreadNotifyRoutineNonSystemCount = 0
		
	print "\nIterating over the nt!PspCreateThreadNotifyRoutine array at {:#x}.".format(PspCreateThreadNotifyRoutine)
	print "Expecting {} nt!PspCreateThreadNotifyRoutineCount and {} nt!PspCreateThreadNotifyRoutineNonSystemCount entries.".format(PspCreateThreadNotifyRoutineCount, PspCreateThreadNotifyRoutineNonSystemCount)
	print spacer

	enumerate_over_callbacks(PspCreateThreadNotifyRoutine)

	print spacer

#####################################################################
# For image callbacks.
#####################################################################
if args.image_callbacks:
	PspLoadImageNotifyRoutine = get_address_from_symbol(mod, "PspLoadImageNotifyRoutine")
	PspLoadImageNotifyRoutineCount = read_dword(get_address_from_symbol(mod, "PspLoadImageNotifyRoutineCount"))

	print "\nIterating over the nt!PspLoadImageNotifyRoutine array at {:#x}.".format(PspLoadImageNotifyRoutine)
	print "Expecting {} nt!PspLoadImageNotifyRoutineCount entries.".format(PspLoadImageNotifyRoutineCount)
	print spacer

	enumerate_over_callbacks(PspLoadImageNotifyRoutine)

	print spacer