Z_createBgmContextCallbackVtable.py

# A helper script to populate vtable in CosCgs of Cos
# @author Nyan Cat
# @category A_Red
# @keybinding
# @menupath
# @toolbar

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

from java.util import ArrayList

from ghidra.program.model.data import CategoryPath, FunctionDefinitionDataType, PointerDataType
from ghidra.program.model.data import DataTypeConflictHandler


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]


funcManager = currentProgram.getFunctionManager()
typeManager = currentProgram.getDataTypeManager()

funcNames = ['_bgm_query_information_service', '_bgm_vbios_table_service', '_bgm_asic_control_service', '_bgm_event_notify', '_bgm_query_hwip_specific_info', '_bgm_update_umd_stable_pstate']
done = [False] * len(funcNames)
ptrs = [None] * len(funcNames)

for func in funcManager.getFunctions(True):
    fullName = str(func)
    funcName = fullName.split("::")[-1]
    if funcName not in funcNames:
        continue

    funcSign = func.getSignature(False)

    funcType = FunctionDefinitionDataType(CategoryPath("/AMDGen/FuncSigns"), funcName + "_sign", funcSign)
    funcType = typeManager.addDataType(funcType, DataTypeConflictHandler.REPLACE_HANDLER)
    ptrType = PointerDataType(funcType, typeManager)
    ptrType = typeManager.addDataType(ptrType, DataTypeConflictHandler.REPLACE_HANDLER)

    for i in range(len(funcNames)):
        if funcNames[i] == funcName:
            assert not done[i]
            done[i] = True
            ptrs[i] = ptrType
            funcNames[i] = fullName

assert all(done[i] or funcNames[i] is None for i in range(len(funcNames)))

serviceCallback = getDataType("BgmContext", typeManager)
for i in range(len(funcNames)):
    if funcNames[i] is None:
        continue
    serviceCallback.replaceAtOffset(8 * (i + 6), ptrs[i], 8, funcNames[i], "Generated by Z_createBgmContextCallbackVtable.py")