GP-1672 - PDB developer scripts for doing basic queries outside of a program.

Ghidra/Features/PDB/developer_scripts/pdbquery/PdbFactory.java

@@ -0,0 +1,242 @@
+/* ###
+ * 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.
+ */
+package pdbquery;
+
+import java.io.IOException;
+import java.util.*;
+import java.util.Map.Entry;
+
+import ghidra.app.script.GhidraScript;
+import ghidra.app.util.bin.format.pdb2.pdbreader.*;
+import ghidra.util.Msg;
+import ghidra.util.exception.CancelledException;
+import ghidra.util.task.TaskMonitor;
+
+/**
+ * Helper class for the PdbQuery set of scripts that allows the PdbQuery scripts to manage the
+ * opening, holding open (caching), and closing of PDB files so that multiple user queries can
+ * be run against the PDBs.  Without this notion, if a user wanted to query a PDB with multiple
+ * runs of a query scripts, each run would have to, once again, open and parse the PDB file
+ * that is desired, which could take minutes or longer.
+ */
+public class PdbFactory {
+
+	private static TreeMap<String, PdbInfo> pdbInfoByFile = new TreeMap<>();
+	private static Map<Class<? extends GhidraScript>, PdbInfo> pdbInfoByScriptClass =
+		new HashMap<>();
+
+	/**
+	 * Opens and retains reference to the PDB file specified.  Must call
+	 * {@link #closePdb(GhidraScript, String)} to close the PDB and remove it from the map; can
+	 * alternatively use {@link #closeAllPdbs(GhidraScript)} to close all PDBs in the map.
+	 * @param script script for which we are working
+	 * @param filename name of PDB file to open and load into map
+	 * @param monitor task monitor
+	 * @return PDB associated with the filename
+	 * @throws CancelledException upon user cancellation
+	 * @throws PdbException upon issues parsing the PDB
+	 */
+	public static PdbInfo openPdb(GhidraScript script, String filename, TaskMonitor monitor)
+			throws CancelledException, PdbException {
+		PdbInfo pdbInfo = pdbInfoByFile.get(filename);
+		if (pdbInfo != null) {
+			return pdbInfo;
+		}
+
+		println(script, "Opening PDB: " + filename);
+
+		try {
+			AbstractPdb pdb = PdbParser.parse(filename, new PdbReaderOptions(), monitor);
+			PdbIdentifiers identifiers = pdb.getIdentifiers();
+			pdb.deserialize(monitor);
+			PdbReaderMetrics metrics = pdb.getPdbReaderMetrics();
+			pdbInfo = new PdbInfo(filename, identifiers, pdb, metrics);
+			pdbInfoByFile.put(filename, pdbInfo);
+			println(script, "\n" + metrics.getPostProcessingReport());
+			return pdbInfo;
+		}
+		catch (IOException ioe) {
+			println(script, ioe.getMessage());
+			Msg.debug(null, ioe.getMessage());
+		}
+		return null;
+	}
+
+	/**
+	 * Closes and unloads the PDB file from the map.  Not removed from map if IOException.
+	 * @param script script for which we are working
+	 * @param filename filename of the PDB file
+	 * @return true if successfully closed and removed from the map; false if not found in the map
+	 * or if problem closing the PDB.
+	 */
+	public static boolean closePdb(GhidraScript script, String filename) {
+		boolean success = closePdbInternal(script, filename);
+		if (success) {
+			pdbInfoByFile.remove(filename);
+		}
+		return success;
+	}
+
+	/**
+	 * Closes and unloads the PDB file from the map.  Not removed from map if IOException.
+	 * @param script script for which we are working
+	 * @return true if all PDBs were successfully closed and unloaded from the map
+	 */
+	public static boolean closeAllPdbs(GhidraScript script) {
+		boolean allUnloaded = true;
+		Iterator<Entry<String, PdbInfo>> iterator = pdbInfoByFile.entrySet().iterator();
+		while (iterator.hasNext()) {
+			Entry<String, PdbInfo> entry = iterator.next();
+			String filename = entry.getKey();
+			boolean success = closePdbInternal(script, filename);
+			if (success) {
+				iterator.remove();
+			}
+			allUnloaded &= success;
+		}
+		return allUnloaded;
+	}
+
+	private static boolean closePdbInternal(GhidraScript script, String filename) {
+		PdbInfo pdbInfo = pdbInfoByFile.get(filename);
+		AbstractPdb pdb = pdbInfo.getPdb();
+		if (pdb != null) {
+			try {
+				pdb.close();
+				String message = "PDB Closed: " + filename;
+				println(script, message);
+			}
+			catch (IOException ioe) {
+				println(script, ioe.getMessage());
+				Msg.info(null, ioe.getMessage());
+				return false;
+			}
+			return true;
+		}
+		return false;
+	}
+
+	/**
+	 * Returns list of PDB information in alphabetical order by filename.
+	 * @return the list
+	 */
+	public static List<PdbInfo> getPdbInfo() {
+		List<PdbInfo> orderedPdbInfo = new ArrayList<>();
+		for (String name : pdbInfoByFile.navigableKeySet()) {
+			orderedPdbInfo.add(pdbInfoByFile.get(name));
+		}
+		return orderedPdbInfo;
+	}
+
+	/**
+	 * Sets the cache PdbInfo value for the class argument
+	 * @param clazz the class for which to cache the value
+	 * @param pdbInfo the PdbInfo value to cache.
+	 */
+	public static void setLastPdbInfoByScriptClass(Class<? extends GhidraScript> clazz,
+			PdbInfo pdbInfo) {
+		pdbInfoByScriptClass.put(clazz, pdbInfo);
+	}
+
+	/**
+	 * Returns the PdbInfo cached for the class argument.
+	 * @param clazz the class of the script used to look up the cached value for return
+	 * @return the PdbInfo
+	 */
+	public static PdbInfo getLastPdbInfoByScriptClass(Class<? extends GhidraScript> clazz) {
+		return pdbInfoByScriptClass.get(clazz);
+	}
+
+	/**
+	 * Method for outputting a message to the console (if script is not null); otherwise outputs
+	 * the message to Msg.info().
+	 * @param script the script
+	 * @param message the message to output to the console
+	 */
+	private static void println(GhidraScript script, String message) {
+		if (script != null) {
+			script.println(message);
+		}
+		else {
+			Msg.info(PdbFactory.class, message);
+		}
+	}
+
+	/**
+	 * Information about a PDB used for specifying and uniquely identifying a PDB along with the
+	 * parsed PDB itself and the PDB parsing metrics generated during the parse.
+	 */
+	public static class PdbInfo {
+		private String filename; // absolute pathname
+		private PdbIdentifiers identifiers;
+		private AbstractPdb pdb;
+		private PdbReaderMetrics metrics;
+
+		/**
+		 * Constructor.
+		 * @param filename PDB filename in absolute pathname format
+		 * @param identifiers identifiers used to help identify versions of the PDB
+		 * @param pdb the parsed PDB
+		 * @param metrics the PDB metrics generated when the PDB was opened and parsed
+		 */
+		PdbInfo(String filename, PdbIdentifiers identifiers, AbstractPdb pdb,
+				PdbReaderMetrics metrics) {
+			this.filename = filename;
+			this.identifiers = identifiers;
+			this.pdb = pdb;
+			this.metrics = metrics;
+		}
+
+		/**
+		 * Returns the PDB filename in absolute path format
+		 * @return the filename
+		 */
+		public String getFilename() {
+			return filename;
+		}
+
+		/**
+		 * Returns the parsed PDB
+		 * @return the parsed PDB
+		 */
+		public AbstractPdb getPdb() {
+			return pdb;
+		}
+
+		/**
+		 * Returns PDB identifiers that help specify its version
+		 * @return the identifiers
+		 */
+		public PdbIdentifiers getIdentifiers() {
+			return identifiers;
+		}
+
+		/**
+		 * Returns the metrics generated during PDB parsing
+		 * @return the metrics
+		 */
+		public PdbReaderMetrics getPdbReaderMetrics() {
+			return metrics;
+		}
+
+		@Override
+		// The PDB and parsing metrics are purposefully not included in this output.
+		public String toString() {
+			return filename + "; " + identifiers;
+		}
+
+	}
+}