ghidra-scripts/find_vtable_offset.py at main · openzt1/ghidra-scripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
#TODO write a description for this script
#@author
#@category _NEW_
#@keybinding
#@menupath
#@toolbar
#@runtime Jython


import re
# from ghidra.program.model.lang import OperandType, Register
from ghidra.program.database.data import StructureDB, PointerDB
# from ghidra.app.decompiler import DecompInterface
# from ghidra.util.task import ConsoleTaskMonitor

from ExportASM import GhidraContext

try:
    from typing import TYPE_CHECKING
except ImportError:
    TYPE_CHECKING = False

if TYPE_CHECKING:
    from ghidra.ghidra_builtins import *


def get_function_at_offset(ctx, offset):
    pass
    # function = ctx.function_manager.getFunctionContaining(offset)
    # if function is None:
    #     return None
    # if not is_identified(function.getName()):
    #     return None
    # return function


def get_function(ctx, addr):
    func = ctx.function_manager.getFunctionContaining(toAddr(addr))
    if func is None:
        print("No function found at the current location.")
    elif func.isThunk():
        print("Function is a thunk.")
    elif func.isExternal():
        print("Function is external.")
    else:
        return func
    return None


entity_names = [
    "BFEntity",
    "BFUnit",
    "BFOverlay",
    "ZTUnit",
    "ZTFood",
    "ZTPath",
    "ZTFence",
    "ZTBuilding",
    "ZTAnimal",
    "ZTGuest",
    "ZTScenery",
    "ZTStaff",
    "ZTKeeper",
    "ZTMaint",
    "ZTGuide",
    "ZTHelicopter",
    "ZTAmbient",
    "ZTRubble",
    "ZTTankWall",
    "ZTTankFilter",
]

entity_type_names = [
    "BFEntityType",
    "BFUnitType",
    "BFOverlayType",
    "ZTUnitType",
    "ZTFoodType",
    "ZTPathType",
    "ZTFenceType",
    "ZTBuildingType",
    "ZTAnimalType",
    "ZTGuestType",
    "ZTSceneryType",
    "ZTStaffType",
    "ZTKeeperType",
    "ZTMaintType",
    "ZTGuideType",
    "ZTHelicopterType",
    "ZTAmbientType",
    "ZTRubbleType",
    "ZTTankWallType",
    "ZTTankFilterType",
]

if __name__ == "__main__":
    ctx = GhidraContext(currentProgram, currentLocation, currentProgram.getFunctionManager(), currentProgram.getSymbolTable())
    types_or_entities = askChoice("Find VTable Offsets", "Entity or Entity Types", ["Entity", "Entity Types"], 0)
    if types_or_entities == "Entity":
        class_names = entity_names
    elif types_or_entities == "Entity Types":
        class_names = entity_type_names
    else:
        print("Invalid choice.")
        exit(1)
    addr = askAddress("Vtable Offset", "Enter Vtable offset (e.g. 0x1c):")
    print("Answer: ", addr, addr.getOffset())
    for datatype in currentProgram.getDataTypeManager().getAllDataTypes():
        parts = datatype.getName().split("::")
        # print(parts)
        if len(parts) > 1 and parts[0] in class_names and (parts[1].startswith("vftable") or parts[1].startswith("vtable")) and type(datatype) is StructureDB:
            # print(datatype.getName())
            f = datatype.getComponentAt(addr.getOffset())
            if f is None:
                print("No datatype: " + parts[0])
                continue
            if f.getDataType() is None or type(f.getDataType()) is not PointerDB:
                print("No pointer: " + str(type(f.getDataType())) + " " + parts[0])
                continue
            if f.getDataType().getDataType() is None:
                print("Void ptr?: " + parts[0])
                continue
            function_address = f.getDataType().getDataType().getName().split("_")[-1]
            print(parts[0] + " " + function_address + " " + get_function(ctx, function_address).getName())
            # print(parts[0] + " " + f.getDataType().getName() + " " + f.getDataType().getDataType().getName())
            # print(datatype.getDefinedComponents()
        # elif len(parts) > 1 and parts[0].split("_")[0] not in ["virt", "~cls", "cls"]:
        #     print(parts[0])