GP-339 Added RecoverClassesFromRTTI script and related classes.

Ghidra/Features/Base/ghidra_scripts/GraphClassesScript.java

@@ -0,0 +1,201 @@
+/* ###
+ * IP: GHIDRA
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * 
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+//Script to graph class hierarchies given metadata found in class structure description that
+// was applied using the ExtractClassInfoFromRTTIScript. 
+//@category C++
+import java.util.*;
+
+import ghidra.app.script.GhidraScript;
+import ghidra.app.services.GraphDisplayBroker;
+import ghidra.framework.plugintool.PluginTool;
+import ghidra.program.model.data.*;
+import ghidra.service.graph.*;
+import ghidra.util.exception.CancelledException;
+import ghidra.util.task.TaskMonitor;
+
+public class GraphClassesScript extends GhidraScript {
+
+	List<Structure> classStructures = new ArrayList<Structure>();
+
+	@Override
+	public void run() throws Exception {
+
+		if (currentProgram == null) {
+			println("There is no open program");
+			return;
+		}
+
+		DataTypeManager dataTypeManager = currentProgram.getDataTypeManager();
+
+		String path = new String("ClassDataTypes");
+		CategoryPath dataTypePath = new CategoryPath(CategoryPath.ROOT, path);
+
+		Category category = dataTypeManager.getCategory(dataTypePath);
+		if (category == null) {
+			println(
+				"/ClassDataTypes folder does not exist so there is no class data to process. Please run the ExtractClassInfoFromRTTIScript to generate the necessary information needed to run this script.");
+			return;
+		}
+
+		Category[] subCategories = category.getCategories();
+
+		getClassStructures(subCategories);
+
+		AttributedGraph graph = createGraph();
+		if (graph.getVertexCount() == 0) {
+			println(
+				"There was no metadata in the class structures so a graph could not be created. Please run the ExtractClassInfoFromRTTIScript to generate the necessary information needed to run this script.");
+		}
+		else {
+			showGraph(graph);
+		}
+
+	}
+
+	private void getClassStructures(Category[] categories) throws CancelledException {
+
+		for (Category category : categories) {
+			monitor.checkCanceled();
+			DataType[] dataTypes = category.getDataTypes();
+			for (DataType dataType : dataTypes) {
+				monitor.checkCanceled();
+				if (dataType.getName().equals(category.getName()) &&
+					dataType instanceof Structure) {
+
+					// if the data type name is the same as the folder name then
+					// it is the main class structure
+					Structure classStructure = (Structure) dataType;
+					if (!classStructures.contains(classStructure)) {
+						classStructures.add(classStructure);
+					}
+
+				}
+			}
+
+			Category[] subcategories = category.getCategories();
+
+			if (subcategories.length > 0) {
+				getClassStructures(subcategories);
+			}
+		}
+	}
+
+	private AttributedGraph createGraph() throws CancelledException {
+
+		AttributedGraph g = new AttributedGraph();
+
+		Iterator<Structure> classStructuresIterator = classStructures.iterator();
+		while (classStructuresIterator.hasNext()) {
+
+			monitor.checkCanceled();
+
+			Structure classStructure = classStructuresIterator.next();
+
+			String description = classStructure.getDescription();
+			String mainClassName = getClassName(description);
+
+			if (mainClassName == null) {
+				continue;
+			}
+
+			AttributedVertex classVertex = g.addVertex(mainClassName);
+
+			int numParents = 0;
+			while (description.contains(":")) {
+
+				numParents++;
+
+				int indexOfColon = description.indexOf(":", 0);
+
+				description = description.substring(indexOfColon + 1);
+
+				int endOfBlock = description.indexOf(":", 0);
+				if (endOfBlock == -1) {
+					endOfBlock = description.length();
+				}
+
+				String parentName = description.substring(0, endOfBlock);
+
+				description = description.substring(endOfBlock);
+
+				boolean isVirtualParent = false;
+				if (parentName.contains("virtual")) {
+					isVirtualParent = true;
+				}
+
+				parentName = parentName.replace("virtual", "");
+				parentName = parentName.replace(" ", "");
+
+
+				AttributedVertex parentVertex = g.addVertex(parentName);
+
+				AttributedEdge edge = g.addEdge(parentVertex, classVertex);
+				if (isVirtualParent) {
+					edge.setAttribute("Color", "Orange");
+				}
+				// else leave it default lime green
+			}
+
+			// no parent = blue vertex
+			if (numParents == 0) {
+				classVertex.setAttribute("Color", "Blue");
+			}
+			// single parent = green vertex
+			else if (numParents == 1) {
+				classVertex.setAttribute("Color", "Green");
+			}
+			// multiple parents = red vertex
+			else {
+				classVertex.setAttribute("Color", "Red");
+			}
+		}
+
+		return g;
+	}
+
+	private void showGraph(AttributedGraph graph) throws Exception {
+
+		GraphDisplay display;
+		PluginTool tool = state.getTool();
+		GraphDisplayBroker broker = tool.getService(GraphDisplayBroker.class);
+		GraphDisplayProvider service = broker.getGraphDisplayProvider("Default Graph Display");
+		display = service.getGraphDisplay(false, TaskMonitor.DUMMY);
+		display.setGraph(graph, "test graph", false, TaskMonitor.DUMMY);
+	}
+
+	private String getClassName(String description) {
+
+		// parse description for class hierarchy
+		if (!description.startsWith("class")) {
+			return null;
+		}
+
+		// skip "class " to get overall class
+		description = description.substring(6);
+		int indexOfColon = description.indexOf(":", 0);
+		String mainClassName;
+		if (indexOfColon == -1) {
+			mainClassName = description;
+		}
+		else {
+			mainClassName = description.substring(0, indexOfColon - 1);
+		}
+		mainClassName = mainClassName.replace(" ", "");
+
+		return mainClassName;
+	}
+
+}