NootKernelcacheVtableAnalyzer.py

# Does all kind of cool stuff to make analyzing classes with vtables in kernelcache easier.
# @author Nyan Cat
# @category A_Red
# @keybinding
# @menupath
# @toolbar

# pyright: reportMissingImports=false
# pyright: reportUndefinedVariable=false

import array
import binascii
import traceback

from java.util import ArrayList

# DataTypeEditorManager.java
# CreatePointerAction.java
from ghidra.program.model.data import CategoryPath, StructureDataType, PointerDataType, FunctionDefinitionDataType

# FindReferencesToField.java -> LocationReferencesPlugin.java -> GenericCompositeDataTypeProgramLocation.java -> ReferenceUtils.java
# Java.
from ghidra.program.model.data import DataTypeConflictHandler

verbose = False
overrideStructs = True


def makeByteArr(length):
    return array.array('b', b'\x00' * length)


def readMem(address, length, littleEndian=True):
    memVal = makeByteArr(length)
    mem.getBytes(address, memVal)
    if littleEndian:
        memVal = memVal[::-1]
    return binascii.hexlify(memVal)


# https://github.com/HackOvert/GhidraSnippets
def getAddress(offset):
    return currentProgram.getAddressFactory().getDefaultAddressSpace().getAddress(offset)


def getDataType(typeName, typeManager):
    matches = ArrayList()
    typeManager.findDataTypes(typeName, matches)
    if len(matches) > 1:
        print("Warning: Using the first " + typeName)
    elif len(matches) == 0:
        return None

    return matches[0]


def makePtrTypeName(typeName, ptrLevel):
    return typeName + " *" * ptrLevel


def ensureDataType(typeName, typeManager, ptrLevel, forceStructSize):
    # Because I kept on misusing this function somehow
    assert "*" not in typeName

    ret = getDataType(makePtrTypeName(typeName, ptrLevel), typeManager)
    if ret is not None:
        # Don't touch other types
        if not isinstance(ret, StructureDataType):
            return ret

        if ptrLevel == 0 and forceStructSize is not None and ret.getLength() != forceStructSize:
            if verbose:
                print("    Resizing {}/{} from {} bytes to {} bytes".format(ret.getCategoryPath(), ret.getName(), ret.getLength(), forceStructSize))
            ret.replaceWith(StructureDataType(ret.getCategoryPath(), ret.getName(), forceStructSize))
        return ret

    # Create our own data type
    if ptrLevel == 0:
        targetLen = forceStructSize if forceStructSize is not None else 0
        if verbose:
            print("Creating {}-bytes struct data type {}".format(targetLen, typeName))
        ret = StructureDataType(CategoryPath("/NootGen/Structs"), typeName, targetLen, typeManager)
    else:
        prevType = ensureDataType(typeName, typeManager, ptrLevel - 1, forceStructSize)
        if verbose:
            print("Creating pointer data type " + makePtrTypeName(typeName, ptrLevel))
        ret = PointerDataType(prevType, typeManager)
    ret = typeManager.addDataType(ret, DataTypeConflictHandler.REPLACE_HANDLER)

    return ret


mem = currentProgram.getMemory()
funcManager = currentProgram.getFunctionManager()
typeManager = currentProgram.getDataTypeManager()
symbolTable = currentProgram.getSymbolTable()

print("Finding classes with vtables...")
# [Type name, Vtable address, Class struct]
vtableClasses = []
it = symbolTable.getAllSymbols(False)
while it.hasNext():
    symbol = it.next()
    typeName = symbol.getName()
    if typeName.endswith("_vtable"):
        className = typeName.split("_vtable")[0]
        vtableClasses.append([className, symbol.getAddress(), getDataType(className, typeManager)])
print("Found {} classes with vtables".format(len(vtableClasses)))

print("")
print("Setting up struct basic structure...")
ulong = ensureDataType("ulong", typeManager, 0, None)
ulongPtr = ensureDataType("ulong", typeManager, 1, None)
setCount = 0
for cur in vtableClasses:
    # Don't touch data type with content
    dataType = cur[2]
    if dataType is None:
        print("Creating {} data type".format(cur[0]))
        dataType = cur[2] = ensureDataType(cur[0], typeManager, 0, 0)

    if dataType.isZeroLength():
        # print("    Appending long * in {}".format(dataType.getName()))
        setCount += 1
        dataType.add(ulongPtr, 8, "vtable", "Generated by NootKernelcacheVtableAnalyzer.py")
        # Ensure the DataType in typeManager is updated
        typeManager.addDataType(dataType, DataTypeConflictHandler.REPLACE_HANDLER)
    elif overrideStructs:
        dataType.replaceAtOffset(0, ulongPtr, 8, "vtable", "Generated by NootKernelcacheVtableAnalyzer.py")
        # Ensure the DataType in typeManager is updated
        typeManager.addDataType(dataType, DataTypeConflictHandler.REPLACE_HANDLER)
print("Set structure on {} structs".format(setCount))

print("")
print("Creating and assigning vtable structs...")

for i in range(len(vtableClasses)):
    try:
        print("{}/{}:".format(i + 1, len(vtableClasses)))
        cur = vtableClasses[i]
        name, vtableAddr, dataType = cur
        vtable = []

        print("    Analyzing {} vtable at {}...".format(name, vtableAddr))
        addr = vtableAddr

        while True:
            ptrAddr = getAddress(readMem(addr, 8))
            if ptrAddr.getOffset() == 0:
                break

            func = funcManager.getFunctionContaining(ptrAddr)
            if func is not None:
                funcName = str(func)
                funcSign = func.getSignature(False)
                vtable.append(funcName)

                funcType = FunctionDefinitionDataType(CategoryPath("/NootGen/FuncSigns"), funcName + "_sign", funcSign)
                typeManager.addDataType(funcType, DataTypeConflictHandler.REPLACE_HANDLER)
            else:
                vtable.append(None)

            addr = addr.add(8)

        if len(vtable) == 0:
            print("    Warning: Failed to find vtable for {}".format(name))
            continue

        print("    Creating {}_vtable with {} entries...".format(name, len(vtable)))
        vtableStruct = ensureDataType(name + "_vtable", typeManager, 0, 0)
        vtableStruct.deleteAll()
        for i in range(len(vtable)):
            funcName = vtable[i]
            if funcName is None:
                funcName = "field_" + str(i).zfill(3)
                vtableStruct.add(ulongPtr, 8, funcName, "Generated by NootKernelcacheVtableAnalyzer.py")
            else:
                signPtr = ensureDataType(funcName + "_sign", typeManager, 1, None)
                vtableStruct.add(signPtr, 8, funcName, "Generated by NootKernelcacheVtableAnalyzer.pyy")
        # Ensure the DataType in typeManager is updated
        typeManager.addDataType(vtableStruct, DataTypeConflictHandler.REPLACE_HANDLER)

        vtableStructPtr = ensureDataType(name + "_vtable", typeManager, 1, None)
        dataType.replaceAtOffset(0, vtableStructPtr, 8, "vtable", "Generated by NootKernelcacheVtableAnalyzer.py")
    except Exception:
        traceback.print_exc()
        print("Ladies and Gentlemen, we fucked up.")