GP-1601 Fixed the ApplyClassFunctionDefinitionUpdatesScript and the ApplyClassFunctionSignatureUpdatesScript and related methods in RecoveredClassHelper to work with new 10.1 function definition changes.

Ghidra/Features/Decompiler/ghidra_scripts/ApplyClassFunctionSignatureUpdatesScript.java

@@ -14,26 +14,28 @@
  * limitations under the License.
  */
 // Script to apply any changes the user has made to recovered class virtual function signatures 
-// edited in the listing. To run the script, put the cursor on any member of the desired class in
-// the listing then run the script. For each function signature in the given class that differs from
-// the associated function definition in the data type manager, the script will update the listing 
-// function signatures of any related virtual functions belonging to parents and children classes.
-// It will also update related data types including function definitions and vftable structures. 
+// edited in the listing. To run the script, put the cursor on a changed virtual function in 
+// the listing then run the script. If the function signature in the given class differs from
+// the associated function definition in the data type manager, the script will update the associated 
+// function definition and any other related function signatures in the listing.
 // Note: The script will not work if the vftable structures were not originally applied to 
 // the vftables using the RecoverClassesFromRTTIScript.
 // At some point, the Ghidra API will be updated to do this automatically instead of needing the 
-// script to do so. 
+// script to do so. For now, to make it a bit easier, you can use the below listed key binding
+// or menupath if you have the "In Tool" checkbox checked for this script in the script manager.
 //@category C++
+//@menupath Scripts.ApplyClassFunctionSignatures
+//@keybinding shift S
 
 import java.util.List;
 
 import classrecovery.RecoveredClassHelper;
 import ghidra.app.script.GhidraScript;
-import ghidra.program.model.data.FunctionDefinition;
-import ghidra.program.model.data.Structure;
+import ghidra.program.model.address.Address;
+import ghidra.program.model.data.*;
+import ghidra.program.model.listing.Data;
 import ghidra.program.model.listing.Function;
 import ghidra.program.model.symbol.Namespace;
-import ghidra.program.model.symbol.Symbol;
 
 public class ApplyClassFunctionSignatureUpdatesScript extends GhidraScript {
 	@Override
@@ -46,29 +48,80 @@ public class ApplyClassFunctionSignatureUpdatesScript extends GhidraScript {
 
 		RecoveredClassHelper classHelper = new RecoveredClassHelper(currentProgram, currentLocation,
 			state.getTool(), this, false, false, false, false, monitor);
+		
+		if(currentAddress == null) {
+			println("Cursor must be in a class function.");
+			return;
+		}
+		Function function = getFunctionContaining(currentAddress);
+		if(function == null) {
+			println("Cursor must be in a class function.");
+			return;
+		}
+		
+		if(function.isThunk()) {
+			println("User should not edit thunks as they are auto-updated from thunked function. " +
+				"Please undo changes to thunk then edit thunked function and rerun script");
+			return;
+		}
+
+		if (function.getName().contains("purecall")) {
+			println("Function definitions are not affected by purecall changes.");
+			return;
+		}
+	
 
 		Namespace classNamespace = classHelper.getClassNamespace(currentAddress);
 		if (classNamespace == null) {
+			println("Cursor must be in a class function.");
+			return;
+		}
+		
+		// get a vftable that points to this function - doesn't matter which since it will
+		// be used to get the underlying function definition which will then be used to update
+		// all related function signatures
+		List<Address> vftablesContainingFunction = classHelper.getVftablesContaining(function);
+
+		// get all vftables that point to given function
+		if (vftablesContainingFunction.isEmpty()) {
 			println(
-				"Either cannot retrieve class namespace or cursor is not in a member of a class namepace");
+				"Function is not a virtual function so has no function definition or related " +
+					"function signatures to update");
 			return;
 		}
 
-		List<Symbol> classVftableSymbols = classHelper.getClassVftableSymbols(classNamespace);
-		if (classVftableSymbols.isEmpty()) {
-			println("There are no vftables in this class");
+		// get one that has a class vftableStructure applied there
+		Address vftableWithAppliedStructure = null;
+		for (Address vftableAddress : vftablesContainingFunction) {
+			monitor.checkCanceled();
+
+			Data dataAt = getDataAt(vftableAddress);
+			if (dataAt == null) {
+				continue;
+			}
+
+			DataType baseDataType = dataAt.getBaseDataType();
+
+			if (baseDataType.getCategoryPath().getPath().contains("ClassDataTypes")) {
+				vftableWithAppliedStructure = vftableAddress;
+				break;
+			}
+		}
+	
+		if (vftableWithAppliedStructure == null) {
+			println(
+				"The vftable(s) containing this function do not have a valid vftable structure " +
+					"applied. Please run the RecoverClassesFromRTTIScript.java on this program before " +
+					"using this script to update virtual functions.");
 			return;
 		}
-
-		println("Applying differing virtual function signatures for class " +
-			classNamespace.getName(true));
-
 		List<Object> changedItems =
-			classHelper.applyNewFunctionSignatures(classNamespace, classVftableSymbols);
+			classHelper.applyNewFunctionSignature(function, vftableWithAppliedStructure);
 
 		if (changedItems.isEmpty()) {
-			println("No differences found for class " + classNamespace.getName(true) +
-				" between the listing virtual function signatures and their associated data type manager function definition data types.");
+			println("No differences found between function signature at " +
+				function.getEntryPoint().toString() +
+				" and its associated function definition in the data type manager.");
 			return;
 		}
 
@@ -77,27 +130,34 @@ public class ApplyClassFunctionSignatureUpdatesScript extends GhidraScript {
 			classHelper.getFunctionDefinitionsOnList(changedItems);
 		List<Function> functionsOnList = classHelper.getFunctionsOnList(changedItems);
 
-		println();
-		println("Updated structures:");
-		for (Structure structure : structuresOnList) {
-			monitor.checkCanceled();
-			println(structure.getPathName());
+		if (!structuresOnList.isEmpty()) {
+			println();
+			println("Updated structures:");
+			for (Structure structure : structuresOnList) {
+				monitor.checkCanceled();
+				println(structure.getPathName());
+			}
 		}
 
-		println();
-		println("Updated function definitions:");
-		for (FunctionDefinition functionDef : functionDefinitionsOnList) {
-			monitor.checkCanceled();
-			println(functionDef.getPathName());
+		if (!functionDefinitionsOnList.isEmpty()) {
+			println();
+			println("Updated function definition:");
+			for (FunctionDefinition functionDef : functionDefinitionsOnList) {
+				monitor.checkCanceled();
+				println(functionDef.getPathName());
+			}
 		}
 
-		println();
-		println("Updated functions:");
-		for (Function function : functionsOnList) {
-			monitor.checkCanceled();
-			println(function.getEntryPoint().toString());
+		if (!functionsOnList.isEmpty()) {
+			println();
+			println("Updated functions:");
+			for (Function functionOnList : functionsOnList) {
+				monitor.checkCanceled();
+				println(functionOnList.getEntryPoint().toString());
+			}
 		}
 
 	}
 
+
 }