Add class fixup Ghidra script

ghidra/ghidra_scripts/ClassFixup.java

@@ -0,0 +1,106 @@
+// Creates classes out namespaces with matching structs, and if they have a
+// vtable, sets the calling convention of the contained function typdefs to
+// __thiscall.
+//
+// For vtables to be found, they must be defined as structs with names ending
+// in "Vtbl" and be pointed to by the first member of a class struct.
+//
+// @category Data Types
+
+import ghidra.app.script.GhidraScript;
+import ghidra.app.util.NamespaceUtils;
+import ghidra.program.database.data.DataTypeUtilities;
+import ghidra.program.model.data.DataTypeComponent;
+import ghidra.program.model.data.FunctionDefinition;
+import ghidra.program.model.data.Pointer;
+import ghidra.program.model.data.Structure;
+import ghidra.program.model.listing.GhidraClass;
+import ghidra.program.model.symbol.Namespace;
+import ghidra.program.model.symbol.Symbol;
+
+
+public class ClassFixup extends GhidraScript {
+	@Override
+	public void run() throws Exception {
+		fixInNamespace(currentProgram.getGlobalNamespace());
+	}
+
+	private void fixInNamespace(final Namespace parent) throws Exception {
+		for (final Symbol s : currentProgram.getSymbolTable()
+		                                    .getChildren(parent.getSymbol()))
+		if (s.getObject() instanceof Namespace ns) switch (ns.getType()) {
+			case Namespace.Type.NAMESPACE:
+				if (shouldBeClass(ns)) {
+					println("Converting \"" + ns.getName(true) + "\" to class...");
+					NamespaceUtils.convertNamespaceToClass(ns);
+
+					// Re-fetch namespace to get its new class version
+					ns = getNamespace(parent, ns.getName());
+				} else {
+					fixInNamespace(ns);
+					break;
+				}
+			case Namespace.Type.CLASS: // fallthrough
+				// Any ns that makes it here should be a class
+				// (can't bind via pattern match because of the
+				// fallthrough)
+				final GhidraClass cls = (GhidraClass)ns;
+
+				// Fix up methods if first member is a vtable
+				if (
+					DataTypeUtilities.findExistingClassStruct(
+						currentProgram.getDataTypeManager(),
+						cls
+					)                        instanceof Structure         struct      &&
+				        struct.getDataTypeAt(0)  instanceof DataTypeComponent memberType  &&
+				        memberType.getDataType() instanceof Pointer           ptrType     &&
+					ptrType   .getDataType() instanceof Structure         vtblT_maybe &&
+					vtblT_maybe.getName().endsWith("Vtbl")
+				) fixMethods(vtblT_maybe);
+				break;
+
+			default: continue;
+		}
+	}
+
+	private boolean shouldBeClass(final Namespace ns) throws Exception {
+	/* Determine if a namespace should be converted to a class
+	The heuristic is whether Ghidra can find a struct that it would be
+	linked with if it was a class.  The process to do so is admittedly a
+	bit byzantine.
+	*/
+		// Change name so the dummy class won't have a conflict
+		final String name = ns.getName();
+		ns.getSymbol()
+		  .setName("__" + ns.getName(), ns.getSymbol().getSource());
+
+		// Create a dummy class to check for a matching struct without
+		// converting the namespace (as classes can't be reverted)
+		final GhidraClass cls = createClass(ns.getParentNamespace(), name);
+
+		// Record whether Ghidra could find a matching struct
+		final boolean ret = DataTypeUtilities.findExistingClassStruct(
+			currentProgram.getDataTypeManager(),
+			cls
+		) != null;
+
+		// Clean up
+		cls.getSymbol().delete();
+		ns.getSymbol().setName(name, ns.getSymbol().getSource());
+
+		return ret;
+	}
+
+	private void fixMethods(final Structure vtbl) throws Exception {
+	/* Set all the calling conventions in a vtable to __thiscall */
+		for (final DataTypeComponent member : vtbl.getComponents())
+		if (
+			member.getDataType() instanceof Pointer            ptr    &&
+			ptr   .getDataType() instanceof FunctionDefinition method &&
+			!method.getCallingConventionName().equals("__thiscall")
+		) {
+			println("Fixing calling convention of \"" + method.getDataTypePath().toString() + "\"...");
+			method.setCallingConvention("__thiscall");
+		}
+	}
+}