mirror of
https://github.com/advanced-threat-research/GhidraScripts
synced 2026-06-08 13:03:27 +00:00
Added the Automagic, GhidrAI, and ColourseByComplexity scripts
This commit is contained in:
ThisIsLibra
+766
@@ -0,0 +1,766 @@
|
||||
//This script is meant to automate the usage of temporary FIDBs, allows you to use multiple BSim databases per file to recover functions, use file metadata recovery scripts (Golang or Nim, for example), and allows you to use a LLM to further annotate the code. Next, the script adds some graphical elements by colourising complex function calls as dark red while non-complex functions are marked as light red.
|
||||
//@author Max 'Libra' Kersten for Trellix
|
||||
//@category
|
||||
//@keybinding
|
||||
//@menupath
|
||||
//@toolbar
|
||||
|
||||
import java.io.File;
|
||||
import java.io.IOException;
|
||||
import java.nio.file.Files;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import ghidra.app.plugin.core.analysis.AutoAnalysisManager;
|
||||
import ghidra.app.script.GhidraScript;
|
||||
import ghidra.app.services.Analyzer;
|
||||
import ghidra.feature.fid.db.FidFile;
|
||||
import ghidra.feature.fid.db.FidFileManager;
|
||||
import ghidra.program.model.listing.Program;
|
||||
import ghidra.util.SystemUtilities;
|
||||
import ghidra.util.exception.CancelledException;
|
||||
|
||||
public class Automagic extends GhidraScript {
|
||||
|
||||
/**
|
||||
* The Function ID analyzer name, taken from Ghidra's source code
|
||||
*/
|
||||
private static final String FUNCTION_ID_ANALYZER = "Function ID";
|
||||
|
||||
/**
|
||||
* The auto analysis manager within Ghidra, used to check if one or more
|
||||
* analyzers are still running
|
||||
*/
|
||||
private AutoAnalysisManager autoAnalysisManager;
|
||||
|
||||
/**
|
||||
* The Function ID database file manager
|
||||
*/
|
||||
private FidFileManager fidFileManager;
|
||||
|
||||
@Override
|
||||
protected void run() throws Exception {
|
||||
/*
|
||||
* Initialise variables which are used in functions later on
|
||||
*/
|
||||
autoAnalysisManager = AutoAnalysisManager.getAnalysisManager(currentProgram);
|
||||
fidFileManager = FidFileManager.getInstance();
|
||||
|
||||
/*
|
||||
* Obtain all required arguments for all scripts
|
||||
*/
|
||||
|
||||
/*
|
||||
* Example CLI input for this script:
|
||||
*
|
||||
* [true/false if extra fidb databases should be used] [path/to/fidb/files]
|
||||
* [true/false if bsim renaming should be used] [path/to/bsim.config]
|
||||
* [true/false if an LLM should be used to rename information] [LLM AI API URL]
|
||||
* [true/false to decide if functions need to be renamed] [true/false to decide
|
||||
* if variables need to be renamed]
|
||||
*/
|
||||
boolean useFidb = askYesNo("Additional FIDB files", "Use additional FIDB files?");
|
||||
String fidFolderPath = askString("Location of additional FIDBs",
|
||||
"What is the folder where the additional FIDB files are located?");
|
||||
boolean useBSimRenamer = askYesNo("Use BSim?", "Use BSim to rename matches?");
|
||||
String bsimConfigPath = askString("BSim configuration file location",
|
||||
"What is the location of the BSim configuration file?");
|
||||
boolean useGhidrAI = askYesNo("GhidrAI LLM usage",
|
||||
"Use the GhidrAI script to rename functions, variables, and summarise functions?");
|
||||
String ghidraiApiUrl = askString("GhidrAI LLM API endpoint", "What is the API URL of the LLM?");
|
||||
boolean ghidraiRenameFunctions = askYesNo("Rename functions", "Rename functions based on AI suggestions?");
|
||||
boolean ghidraiRenameVariables = askYesNo("Rename variables", "Rename variables based on AI suggestions?");
|
||||
|
||||
/*
|
||||
* End of obtaining arguments for all scripts
|
||||
*/
|
||||
|
||||
/*
|
||||
* Test the validity of all script arguments, prior to making changes to the
|
||||
* current program. Values asked via the askValues (or specific ask* functions)
|
||||
* cannot be null, so there is no need to check for the null-state.
|
||||
*/
|
||||
|
||||
File fidFolder = new File(fidFolderPath);
|
||||
if (useFidb) {
|
||||
// Check if the FID folder exists
|
||||
if (fidFolder.exists() == false) {
|
||||
throw new IOException("The provided FIDB folder does not exist: " + fidFolder.getAbsolutePath());
|
||||
}
|
||||
|
||||
/*
|
||||
* Check if the provided existing path is a folder by throwing an exception if
|
||||
* it is a file
|
||||
*/
|
||||
if (fidFolder.isFile()) {
|
||||
throw new IOException("The provided FIDB folder is a file: " + fidFolder.getAbsolutePath());
|
||||
}
|
||||
}
|
||||
|
||||
List<BSimExecutionConfig> configs = null;
|
||||
if (useBSimRenamer) {
|
||||
File bsimConfigFile = new File(bsimConfigPath);
|
||||
|
||||
// Check if the BSim config file exists
|
||||
if (bsimConfigFile.exists() == false) {
|
||||
throw new IOException("The provided BSim file does not exist: " + bsimConfigPath);
|
||||
}
|
||||
|
||||
/*
|
||||
* Check if the provided existing path is a file by throwing an exception if it
|
||||
* is a folder
|
||||
*/
|
||||
if (bsimConfigFile.isDirectory()) {
|
||||
throw new IOException("The provided BSim path is a directory: " + bsimConfigPath);
|
||||
}
|
||||
|
||||
// Parse all configs from the existing BSim config file
|
||||
configs = parseBSimExecutionConfigs(bsimConfigFile);
|
||||
}
|
||||
|
||||
/*
|
||||
* Metadata analysis is the closest ground truth that is available from the
|
||||
* lossy compilation process. Recovering names from Golang's pclntab and
|
||||
* demangling Nim function names are included, but any type of metadata related
|
||||
* script should be put in this section.
|
||||
*/
|
||||
|
||||
/*
|
||||
* Recover function names, types, and strings (both dynamic and static), based
|
||||
* on the pclntab. More information can be found here:
|
||||
* https://www.trellix.com/blogs/research/feeding-gophers-to-ghidra/
|
||||
*/
|
||||
runScriptAndWait("GolangRecovery", false);
|
||||
|
||||
/*
|
||||
* Demangle Nim function names and recover Nim strings, made by ESET's Alexandre
|
||||
* Côté Cyr. More information on ESET's blog:
|
||||
* https://www.welivesecurity.com/en/eset-research/introducing-nimfilt-reverse-
|
||||
* engineering-tool-nim-compiled-binaries/
|
||||
*
|
||||
* The GitHub repository with the scripts: https://github.com/eset/nimfilt
|
||||
*/
|
||||
runScriptAndWait("NimFilt", false);
|
||||
|
||||
/*
|
||||
* Load and activate all FID databases from a given folder. All newly loaded and
|
||||
* activated FID databases are returned in a list, allowing for easy removal
|
||||
* from the attached FID databases list later on.
|
||||
*/
|
||||
List<FidFile> additionalFids = loadAdditionalFunctionIdDatabases(fidFolder);
|
||||
|
||||
// Apply FunctionID signatures from the loaded FID databases
|
||||
runAnalyzer(FUNCTION_ID_ANALYZER);
|
||||
|
||||
// Run the BSim rename script for all given and previously parsed configs
|
||||
if (useBSimRenamer) {
|
||||
runBSimRenamer(configs);
|
||||
}
|
||||
|
||||
// Run the GhidrAI rename and summary script
|
||||
if (useGhidrAI) {
|
||||
runGhidrAI(ghidraiApiUrl, ghidraiRenameFunctions, ghidraiRenameVariables);
|
||||
}
|
||||
|
||||
/*
|
||||
* Propagate the parameters for external functions. This is done by running
|
||||
* Ghidra's Propagatex86ExternalParams for 32-bit programs. For 64-bit programs,
|
||||
* the variant of Karsten Hahn (aka struppigel) is used (originally found at
|
||||
* https://github.com/struppigel/hedgehog-tools/blob/main/ghidra_scripts/
|
||||
* PropagateExternalParametersX64.java)
|
||||
*
|
||||
* The function to get the bitness returns -1 if an error is found, which is why
|
||||
* the if- and else-if-clauses are set to a specific number without the usage of
|
||||
* an else-clause.
|
||||
*/
|
||||
int bitness = getBitness();
|
||||
if (bitness == 32) {
|
||||
runScriptAndWait("PropagateExternalParametersScript", false);
|
||||
} else if (bitness == 64) {
|
||||
runScriptAndWait("PropagateExternalParametersX64", false);
|
||||
}
|
||||
|
||||
/*
|
||||
* Colour function call instructions based on the complexity depth level of the
|
||||
* called function
|
||||
*/
|
||||
runScriptAndWait("ColouriseByComplexity", false);
|
||||
|
||||
// Restore original FID database selection
|
||||
restoreOriginalFunctionIdDatabaseSelection(additionalFids);
|
||||
}
|
||||
|
||||
/**
|
||||
* Prints the given message via the
|
||||
* {@link ghidra.app.script.GhidraScript#println(String)} method, depending on
|
||||
* the running mode and boolean.
|
||||
*
|
||||
* @param message the message to print
|
||||
* @param printHeadless true if the message should be printed when executing
|
||||
* headless, false if not
|
||||
*/
|
||||
private void log(String message, boolean printHeadless) {
|
||||
if (SystemUtilities.isInHeadlessMode()) {
|
||||
if (printHeadless == false) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
println(message);
|
||||
}
|
||||
|
||||
/**
|
||||
* A wrapper for {@link #runScriptAndWait(String, String[], boolean)} to run a
|
||||
* script without any arguments
|
||||
*
|
||||
* @param scriptName the name of the script
|
||||
* @param throwExceptions if the execution of the script should throw
|
||||
* encountered exceptions, or if they should be caught
|
||||
* @throws Exception the exception which is thrown if the second argument is
|
||||
* true
|
||||
*/
|
||||
private void runScriptAndWait(String scriptName, boolean throwExceptions) throws Exception {
|
||||
runScriptAndWait(scriptName, null, throwExceptions);
|
||||
}
|
||||
|
||||
/**
|
||||
* Runs a script and waits until its execution is complete
|
||||
*
|
||||
* @param scriptName the name of the script
|
||||
* @param scriptArguments arguments for the script, may be <code>null</code>
|
||||
* @param throwExceptions if the execution of the script should throw
|
||||
* encountered exceptions, or if they should be caught
|
||||
* @throws Exception
|
||||
*/
|
||||
private void runScriptAndWait(String scriptName, String[] scriptArguments, boolean throwExceptions)
|
||||
throws Exception {
|
||||
try {
|
||||
log("Running " + scriptName, true);
|
||||
runScript(scriptName, scriptArguments, getState());
|
||||
waitUntilAutoAnalysisCompletes();
|
||||
log("Finished " + scriptName, true);
|
||||
} catch (Exception ex) {
|
||||
if (throwExceptions) {
|
||||
throw ex;
|
||||
}
|
||||
printerr(ex.toString());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Run an analyzer based on its name
|
||||
*
|
||||
* @param analyzerName the name of the analyzer to run
|
||||
* @throws InterruptedException
|
||||
* @throws CancelledException
|
||||
*/
|
||||
private void runAnalyzer(String analyzerName) throws InterruptedException, CancelledException {
|
||||
Analyzer analyzer = autoAnalysisManager.getAnalyzer(analyzerName);
|
||||
autoAnalysisManager.scheduleOneTimeAnalysis(analyzer, currentProgram.getAddressFactory().getAddressSet());
|
||||
waitUntilAutoAnalysisCompletes();
|
||||
}
|
||||
|
||||
/**
|
||||
* Wait until no analyzers are running anymore
|
||||
*
|
||||
* @throws InterruptedException
|
||||
* @throws CancelledException
|
||||
*/
|
||||
private void waitUntilAutoAnalysisCompletes() throws InterruptedException, CancelledException {
|
||||
while (autoAnalysisManager.isAnalyzing()) {
|
||||
if (monitor.isCancelled()) {
|
||||
throw new CancelledException();
|
||||
}
|
||||
Thread.sleep(1000); // Sleep to avoid consuming extra CPU cycles in-between checks
|
||||
}
|
||||
}
|
||||
|
||||
private void runGhidrAI(String apiUrl, boolean renameFunctions, boolean renameVariables) throws Exception {
|
||||
/*
|
||||
* Required information to run GhidrAI:
|
||||
*
|
||||
* - AI API URL
|
||||
*
|
||||
* - Rename function based on suggestion: yes/no
|
||||
*
|
||||
* - Rename variables based on suggestions: yes/no
|
||||
*/
|
||||
String[] ghidraiArgs = new String[3];
|
||||
ghidraiArgs[0] = apiUrl; // AI API URL
|
||||
ghidraiArgs[1] = "" + renameFunctions; // suggest function name
|
||||
ghidraiArgs[2] = "" + renameVariables; // rename based on suggestion
|
||||
runScriptAndWait("GhidrAI", ghidraiArgs, false);
|
||||
}
|
||||
|
||||
private String replaceLastOccurrence(String original, String target, String replacement) {
|
||||
int lastIndex = original.lastIndexOf(target);
|
||||
if (lastIndex == -1) {
|
||||
return original; // Target string not found
|
||||
}
|
||||
String beforeLastOccurrence = original.substring(0, lastIndex);
|
||||
String afterLastOccurrence = original.substring(lastIndex + target.length());
|
||||
return beforeLastOccurrence + replacement + afterLastOccurrence;
|
||||
}
|
||||
|
||||
private List<BSimExecutionConfig> parseBSimExecutionConfigs(File file) throws IOException {
|
||||
List<BSimExecutionConfig> configs = new ArrayList<>();
|
||||
int nonParsableLines = 0;
|
||||
|
||||
List<String> lines = new ArrayList<>();
|
||||
try {
|
||||
lines = Files.readAllLines(file.toPath());
|
||||
} catch (IOException ex) {
|
||||
printerr("Failed to read the file with the referenced BSim files: " + file.getAbsolutePath());
|
||||
throw ex;
|
||||
}
|
||||
|
||||
if (lines.isEmpty()) {
|
||||
throw new IOException("No lines present in the BSim config file!");
|
||||
}
|
||||
|
||||
for (String line : lines) {
|
||||
if (line.startsWith("#")) {
|
||||
nonParsableLines++;
|
||||
continue;
|
||||
}
|
||||
int errorCount = 0;
|
||||
println("Parsing the following BSim config file line: " + line);
|
||||
|
||||
String[] split = line.split(",");
|
||||
if (split.length < 6 || split.length > 7) {
|
||||
printerr(
|
||||
"Not all data is present within the given line, the length should be 6 or 7, whereas the length of this line is only "
|
||||
+ split.length);
|
||||
errorCount++;
|
||||
}
|
||||
|
||||
String databaseUrl = split[0];
|
||||
if (databaseUrl.endsWith(".mv.db")) {
|
||||
databaseUrl = replaceLastOccurrence(databaseUrl, ".mv.db", "");
|
||||
println("Removed the .mv.db extension from the database file path: " + databaseUrl);
|
||||
}
|
||||
File database = new File(databaseUrl + ".mv.db");
|
||||
if (database.exists() == false) {
|
||||
printerr("The provided database does not exist: " + databaseUrl);
|
||||
errorCount++;
|
||||
}
|
||||
if (database.isDirectory()) {
|
||||
printerr("The provided database url is a folder: " + databaseUrl);
|
||||
errorCount++;
|
||||
}
|
||||
|
||||
double lowerSimilartyBound = -1;
|
||||
try {
|
||||
lowerSimilartyBound = Double.parseDouble(split[1]);
|
||||
} catch (Exception ex) {
|
||||
printerr("The lower similarity bound is not a double: " + split[1]);
|
||||
errorCount++;
|
||||
}
|
||||
|
||||
int maximumNumberOfMatches = -1;
|
||||
try {
|
||||
maximumNumberOfMatches = Integer.parseInt(split[2]);
|
||||
} catch (Exception ex) {
|
||||
printerr("The maximum number of matches is not an integer: " + split[2]);
|
||||
errorCount++;
|
||||
}
|
||||
|
||||
/*
|
||||
* Booleans parsed by Boolean.parseBoolean (and similar functions) return true
|
||||
* for any spelling of true, and return false for anything else. This is not the
|
||||
* check we want here, as we want to have either true or false as a value,
|
||||
* anything else is incorrect. Thus, testing of the literal string value is what
|
||||
* we need.
|
||||
*/
|
||||
String renameSingleMatchesString = split[3];
|
||||
boolean renameSingleMatches = false;
|
||||
if (renameSingleMatchesString.equalsIgnoreCase("true")) {
|
||||
renameSingleMatches = true;
|
||||
} else if (renameSingleMatchesString.equalsIgnoreCase("false")) {
|
||||
renameSingleMatches = false;
|
||||
} else {
|
||||
printerr("The boolean value to define if single matches should be renamed is neither true nor false: "
|
||||
+ renameSingleMatches);
|
||||
errorCount++;
|
||||
}
|
||||
|
||||
String renameMultiMatchesString = split[4];
|
||||
boolean renameMultiMatches = false;
|
||||
if (renameMultiMatchesString.equalsIgnoreCase("true")) {
|
||||
renameMultiMatches = true;
|
||||
} else if (renameMultiMatchesString.equalsIgnoreCase("false")) {
|
||||
renameMultiMatches = false;
|
||||
} else {
|
||||
printerr("The boolean value to define if multi-matches should be renamed is neither true nor false: "
|
||||
+ renameMultiMatchesString);
|
||||
errorCount++;
|
||||
}
|
||||
|
||||
String renameGenericMatchesString = split[5];
|
||||
boolean renameGenericMatches = false;
|
||||
if (renameGenericMatchesString.equalsIgnoreCase("true")) {
|
||||
renameGenericMatches = true;
|
||||
} else if (renameGenericMatchesString.equalsIgnoreCase("false")) {
|
||||
renameGenericMatches = false;
|
||||
} else {
|
||||
printerr("The boolean value to define if generic matches should be renamed is neither true nor false: "
|
||||
+ renameGenericMatchesString);
|
||||
errorCount++;
|
||||
}
|
||||
|
||||
String customPrefix = null;
|
||||
if (split.length == 7) {
|
||||
customPrefix = split[6];
|
||||
if (customPrefix == null || customPrefix.isBlank()) {
|
||||
printerr("The required custom prefix has not been provided or is left blank while it is needed!");
|
||||
errorCount++;
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
* Required arguments, in order, split by a comma, should be:
|
||||
*
|
||||
* - database URL, without any of the file extensions (the H2 database has two,
|
||||
* neither should be included)
|
||||
*
|
||||
* - the lower similarity bound between 0 and 1, as a double (i.e. 0.7)
|
||||
*
|
||||
* - the maximum number of BSim matches per local function (i.e. 20)
|
||||
*
|
||||
* - a boolean (true or false) which decides if single matches should be renamed
|
||||
*
|
||||
* - a boolean (true or false) which decides if multi-matches should be renamed
|
||||
*
|
||||
* - a boolean (true or false) which decides if generic matches should be
|
||||
* renamed (if true, a 7th argument is required, otherwise it isn't)
|
||||
*
|
||||
* - a String which is the prefix used when renaming generic matches
|
||||
*/
|
||||
|
||||
// If there are no errors, this value should be zero
|
||||
if (errorCount == 0) {
|
||||
// Create the custom object
|
||||
BSimExecutionConfig config = new BSimExecutionConfig(databaseUrl, lowerSimilartyBound,
|
||||
maximumNumberOfMatches, renameSingleMatches, renameMultiMatches, renameGenericMatches,
|
||||
customPrefix);
|
||||
// Add the config to the list
|
||||
configs.add(config);
|
||||
// Notify the analyst
|
||||
println("Created the config object!");
|
||||
} else {
|
||||
// Print the error
|
||||
printerr("The config hasn't been parsed!");
|
||||
}
|
||||
}
|
||||
|
||||
// Check if all lines have been parsed
|
||||
if ((configs.size() + nonParsableLines) == lines.size()) {
|
||||
// All lines have been parsed successfully
|
||||
return configs;
|
||||
}
|
||||
// If this is not the case, there were errors
|
||||
throw new IOException("Not all BSim configs were parsed correctly!");
|
||||
}
|
||||
|
||||
private void runBSimRenamer(List<BSimExecutionConfig> configs) throws Exception {
|
||||
for (BSimExecutionConfig config : configs) {
|
||||
/*
|
||||
* If the database does not match the current program, it is not to be used.
|
||||
* There wont be (valid) results for the current program, and will only consume
|
||||
* a lot of time. Using a large H2 database is never fast, but can be terribly
|
||||
* slow if it is really large. A high(er) number of maximum matches will also
|
||||
* increase the runtime.
|
||||
*
|
||||
* This method is based on the naming scheme of the H2 databases that are
|
||||
* linked, which is not the easiest approach, but works well as a best-effort
|
||||
* method.
|
||||
*/
|
||||
if (bsimDatabaseMatches(config.getDatabaseUrl()) == false) {
|
||||
continue;
|
||||
}
|
||||
|
||||
int size = 6;
|
||||
if (config.getPrefix() != null) {
|
||||
size = 7;
|
||||
}
|
||||
String[] bsimFunctionRenamerArgs = new String[size];
|
||||
bsimFunctionRenamerArgs[0] = config.getDatabaseUrl(); // database URL
|
||||
bsimFunctionRenamerArgs[1] = "" + config.getLowerSimilarityBound(); // lower similarity bound
|
||||
bsimFunctionRenamerArgs[2] = "" + config.getMaximumMatches(); // maximum bsim matches per local function
|
||||
bsimFunctionRenamerArgs[3] = "" + config.renameSingleMatches(); // rename single matches
|
||||
bsimFunctionRenamerArgs[4] = "" + config.renameMultiMatches(); // rename multi-matches
|
||||
bsimFunctionRenamerArgs[5] = "" + config.renameGenericMatches(); // rename generic matches
|
||||
if (size == 7) {
|
||||
bsimFunctionRenamerArgs[6] = config.getPrefix(); // the custom prefix to use
|
||||
}
|
||||
|
||||
runScriptAndWait("BsimFunctionRenamer", bsimFunctionRenamerArgs, false);
|
||||
}
|
||||
}
|
||||
|
||||
private boolean bsimDatabaseMatches(String databaseUrl) {
|
||||
/*
|
||||
* TODO handle bsim files based on their file name, automatically excluding
|
||||
* architectures which are incorrect for the given file, thus saving a lot of
|
||||
* time when dealing with large H2 databases.
|
||||
*
|
||||
* This can be done by running some extra checks:
|
||||
*
|
||||
* -Check the language ID of the BSim database and the current program (keep in
|
||||
* mind that "medium nosize" allows 32-bit and 64-bit signatures to
|
||||
* mix-and-match
|
||||
*
|
||||
* C:\Users\malwa\Desktop\bsim-signatures\bsim.rust.windows.x86_64.h2.medium-
|
||||
* nosize
|
||||
*/
|
||||
|
||||
/*
|
||||
* Assume the database is a file which exists, as this check has been performed
|
||||
* before
|
||||
*/
|
||||
File databaseUrlFile = new File(databaseUrl);
|
||||
BsimFileNameContainer container = convertIntoContainer(databaseUrlFile.getName());
|
||||
// Container is null if the file name format does not match
|
||||
if (container == null) {
|
||||
return true;
|
||||
}
|
||||
|
||||
String executableFormat = currentProgram.getExecutableFormat(); // i.e. Portable Executable (PE)
|
||||
String cspec = currentProgram.getCompilerSpec().getCompilerSpecID().getIdAsString(); // i.e. windows
|
||||
String languageId = currentProgram.getLanguage().getLanguageID().getIdAsString(); // i.e. x86:LE:32:default
|
||||
|
||||
if (container.getPlatform().equalsIgnoreCase("windows")) {
|
||||
if (executableFormat.equalsIgnoreCase("Portable Executable (PE)")
|
||||
|| cspec.equalsIgnoreCase(container.getPlatform())) {
|
||||
// Matches, continue onwards
|
||||
}
|
||||
} else if (container.getPlatform().equalsIgnoreCase("linux")) {
|
||||
if (executableFormat.equalsIgnoreCase("Executable and Linking Format (ELF)")
|
||||
|| cspec.equalsIgnoreCase(container.getPlatform())) {
|
||||
// Matches, continue onwards
|
||||
}
|
||||
} else if (container.getPlatform().equalsIgnoreCase("apple")
|
||||
|| container.getPlatform().equalsIgnoreCase("macos")) {
|
||||
if (executableFormat.equalsIgnoreCase("Mac OS X Mach-O")
|
||||
|| cspec.equalsIgnoreCase(container.getPlatform())) {
|
||||
// Matches, continue onwards
|
||||
}
|
||||
}
|
||||
|
||||
String[] languageIdSplit = languageId.split(":");
|
||||
if (languageIdSplit.length == 4) {
|
||||
// TODO get uniform check for the language ID
|
||||
/*
|
||||
* TODO ensure that the databaseSize cases with "medium-nosize" includes cross
|
||||
* architecture options as 32-bits and 64-bits signatures are used
|
||||
* interchangeably (to a certain extend)
|
||||
*/
|
||||
String architecture = languageIdSplit[0] + languageIdSplit[2];
|
||||
} else {
|
||||
// Cannot confirm the type, but it should have a length of 4: x86:LE:32:default
|
||||
}
|
||||
|
||||
// TODO update the return value once the function is complete
|
||||
return true;
|
||||
}
|
||||
|
||||
private BsimFileNameContainer convertIntoContainer(String fileName) {
|
||||
String[] split = fileName.split("\\.");
|
||||
if (split.length == 6) {
|
||||
String bsim = split[0];
|
||||
String library = split[1];
|
||||
String platform = split[2];
|
||||
String architecture = split[3];
|
||||
String databaseType = split[4];
|
||||
String databaseSize = split[5];
|
||||
return new BsimFileNameContainer(bsim, library, platform, architecture, databaseType, databaseSize);
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
class BsimFileNameContainer {
|
||||
private String bsim;
|
||||
private String library;
|
||||
private String platform;
|
||||
private String architecture;
|
||||
private String databaseType;
|
||||
private String databaseSize;
|
||||
|
||||
public BsimFileNameContainer(String bsim, String library, String platform, String architecture,
|
||||
String databaseType, String databaseSize) {
|
||||
super();
|
||||
this.bsim = bsim;
|
||||
this.library = library;
|
||||
this.platform = platform;
|
||||
this.architecture = architecture;
|
||||
this.databaseType = databaseType;
|
||||
this.databaseSize = databaseSize;
|
||||
}
|
||||
|
||||
public String getBsim() {
|
||||
return bsim;
|
||||
}
|
||||
|
||||
public String getLibrary() {
|
||||
return library;
|
||||
}
|
||||
|
||||
public String getPlatform() {
|
||||
return platform;
|
||||
}
|
||||
|
||||
public String getArchitecture() {
|
||||
return architecture;
|
||||
}
|
||||
|
||||
public String getDatabaseType() {
|
||||
return databaseType;
|
||||
}
|
||||
|
||||
public String getDatabaseSize() {
|
||||
return databaseSize;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Loads all files ending with .FIDB from the givne folder
|
||||
*
|
||||
* @param folder the folder to load the function ID files from
|
||||
* @return a list of loaded FID files
|
||||
*/
|
||||
private List<FidFile> loadAdditionalFunctionIdDatabases(File folder) {
|
||||
List<FidFile> fids = new ArrayList<>();
|
||||
|
||||
for (File file : folder.listFiles()) {
|
||||
if (file.isFile() == false) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (file.getName().endsWith(".fidb")) {
|
||||
int size = fidFileManager.getFidFiles().size();
|
||||
FidFile fidFile = fidFileManager.addUserFidFile(file);
|
||||
if (fidFile != null) {
|
||||
fidFile.setActive(true);
|
||||
/*
|
||||
* Adding a valid FIDB file is always possible, even if its already known to
|
||||
* Ghidra. Thus, a check is added to see if the total number of FIDB files has
|
||||
* increased. If it was already known, this is not the case. If it wasn't, then
|
||||
* it is. To avoid removing a FIDB from the already known FIDB list, the FIDB
|
||||
* file will only be added to the list if Ghidra added a new FIDB to its known
|
||||
* list.
|
||||
*/
|
||||
if (fidFileManager.getFidFiles().size() > size) {
|
||||
fids.add(fidFile);
|
||||
println("Added \"" + fidFile.getName() + "\" to the active loaded FunctionID databases!");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return fids;
|
||||
}
|
||||
|
||||
/**
|
||||
* Restore the original FIDB selection. This unloads the previously loaded FIDBs
|
||||
* @param extraFunctionIdDatabases the list of previously loaded FID files
|
||||
*/
|
||||
private void restoreOriginalFunctionIdDatabaseSelection(List<FidFile> extraFunctionIdDatabases) {
|
||||
for (FidFile fidFile : extraFunctionIdDatabases) {
|
||||
fidFileManager.removeUserFile(fidFile);
|
||||
println("Removed \"" + fidFile.getName() + "\" from the loaded FunctionID databases!");
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets the program's bitness
|
||||
*
|
||||
* @return the bitness as an integer, or -1 if an error occurs
|
||||
*/
|
||||
private int getBitness() {
|
||||
Map<String, String> mapping = currentProgram.getMetadata();
|
||||
String bitness = mapping.get("Address Size");
|
||||
try {
|
||||
return Integer.parseInt(bitness);
|
||||
} catch (Exception ex) {
|
||||
printerr(ex.toString());
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
class BSimExecutionConfig {
|
||||
private String databaseUrl;
|
||||
private double lowerSimilarityBound;
|
||||
private int maximumMatches;
|
||||
private boolean renameSingleMatches;
|
||||
private boolean renameMultiMatches;
|
||||
private boolean renameGenericMatches;
|
||||
private String prefix;
|
||||
|
||||
public BSimExecutionConfig(String databaseUrl, double lowerSimilarityBound, int maximumMatches,
|
||||
boolean renameSingleMatches, boolean renameMultiMatches, boolean renameGenericMatches, String prefix) {
|
||||
this.databaseUrl = databaseUrl;
|
||||
this.lowerSimilarityBound = lowerSimilarityBound;
|
||||
this.maximumMatches = maximumMatches;
|
||||
this.renameSingleMatches = renameSingleMatches;
|
||||
this.renameMultiMatches = renameMultiMatches;
|
||||
this.renameGenericMatches = renameGenericMatches;
|
||||
this.prefix = prefix;
|
||||
}
|
||||
|
||||
public String getDatabaseUrl() {
|
||||
return databaseUrl;
|
||||
}
|
||||
|
||||
public void setDatabaseUrl(String databaseUrl) {
|
||||
this.databaseUrl = databaseUrl;
|
||||
}
|
||||
|
||||
public double getLowerSimilarityBound() {
|
||||
return lowerSimilarityBound;
|
||||
}
|
||||
|
||||
public void setLowerSimilarityBound(double lowerSimilarityBound) {
|
||||
this.lowerSimilarityBound = lowerSimilarityBound;
|
||||
}
|
||||
|
||||
public int getMaximumMatches() {
|
||||
return maximumMatches;
|
||||
}
|
||||
|
||||
public void setMaximumMatches(int maximumMatches) {
|
||||
this.maximumMatches = maximumMatches;
|
||||
}
|
||||
|
||||
public boolean renameSingleMatches() {
|
||||
return renameSingleMatches;
|
||||
}
|
||||
|
||||
public void setRenameSingleMatches(boolean renameSingleMatches) {
|
||||
this.renameSingleMatches = renameSingleMatches;
|
||||
}
|
||||
|
||||
public boolean renameMultiMatches() {
|
||||
return renameMultiMatches;
|
||||
}
|
||||
|
||||
public void setRenameMultiMatches(boolean renameMultiMatches) {
|
||||
this.renameMultiMatches = renameMultiMatches;
|
||||
}
|
||||
|
||||
public boolean renameGenericMatches() {
|
||||
return renameGenericMatches;
|
||||
}
|
||||
|
||||
public void setRenameGenericMatches(boolean renameGenericMatches) {
|
||||
this.renameGenericMatches = renameGenericMatches;
|
||||
}
|
||||
|
||||
public String getPrefix() {
|
||||
return prefix;
|
||||
}
|
||||
|
||||
public void setPrefix(String prefix) {
|
||||
this.prefix = prefix;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,232 @@
|
||||
//Colourises the complexity of a function where it is called in the disassembly listing. Light red is not complex, dark red is complex.
|
||||
//@author Max 'Libra' Kersten for Trellix. The graph related code has been taken from and inspired by Ghidra's base: https://github.com/NationalSecurityAgency/ghidra/blob/master/Ghidra/Features/Base/src/main/java/ghidra/app/cmd/module/ComplexityDepthModularizationCmd.java#L43
|
||||
//@category
|
||||
//@keybinding
|
||||
//@menupath
|
||||
//@toolbar
|
||||
|
||||
import java.awt.Color;
|
||||
import java.util.ArrayList;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import ghidra.app.decompiler.DecompInterface;
|
||||
import ghidra.app.script.GhidraScript;
|
||||
import ghidra.app.util.viewer.listingpanel.PropertyBasedBackgroundColorModel;
|
||||
import ghidra.graph.GDirectedGraph;
|
||||
import ghidra.graph.GraphAlgorithms;
|
||||
import ghidra.graph.GraphFactory;
|
||||
import ghidra.program.database.IntRangeMap;
|
||||
import ghidra.program.model.address.Address;
|
||||
import ghidra.program.model.block.BasicBlockModel;
|
||||
import ghidra.program.model.block.CodeBlock;
|
||||
import ghidra.program.model.block.CodeBlockIterator;
|
||||
import ghidra.program.model.block.CodeBlockReference;
|
||||
import ghidra.program.model.block.CodeBlockReferenceIterator;
|
||||
import ghidra.program.model.block.graph.CodeBlockEdge;
|
||||
import ghidra.program.model.block.graph.CodeBlockVertex;
|
||||
import ghidra.program.model.listing.Function;
|
||||
import ghidra.program.model.symbol.Reference;
|
||||
import ghidra.program.model.symbol.ReferenceIterator;
|
||||
import ghidra.util.exception.CancelledException;
|
||||
import ghidra.util.exception.DuplicateNameException;
|
||||
|
||||
public class ColouriseByComplexity extends GhidraScript {
|
||||
|
||||
@Override
|
||||
protected void run() throws Exception {
|
||||
// Initialise the decompiler
|
||||
DecompInterface decompiler = new DecompInterface();
|
||||
// Open the current program
|
||||
decompiler.openProgram(currentProgram);
|
||||
|
||||
// Get the call graph
|
||||
GDirectedGraph<CodeBlockVertex, CodeBlockEdge> callGraph = createCallGraph();
|
||||
// Obtain the complexity depth graph
|
||||
Map<CodeBlockVertex, Integer> complexityDepth = GraphAlgorithms.getComplexityDepth(callGraph);
|
||||
// Get the list of all levels within the graph
|
||||
List<List<Function>> partition = createFunctionList(complexityDepth);
|
||||
|
||||
// Remove the empty levels from the list
|
||||
partition = cleanList(partition);
|
||||
|
||||
// Iterate over the levels
|
||||
for (int i = 0; partition.size() > i; i++) {
|
||||
// Get the current list (or level, if you will)
|
||||
List<Function> list = partition.get(i);
|
||||
|
||||
// Iterate over the functions within the current list
|
||||
for (Function function : list) {
|
||||
// Exclude thunk and external functions
|
||||
if (function.isThunk() || function.isExternal()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
/*
|
||||
* Store the number of levels as a double, required to perform arithmetics that
|
||||
* return a double
|
||||
*/
|
||||
double size = partition.size();
|
||||
|
||||
// Get the current value, incremented by one since i is zero-based
|
||||
double currentValue = i + 1;
|
||||
|
||||
/*
|
||||
* Calculate the percentage of this level is relatively to the rest all levels
|
||||
*/
|
||||
double percentage = currentValue / size;
|
||||
|
||||
/*
|
||||
* One minus the percentage, since the first entry in the partition list, has
|
||||
* the highest complexity depth
|
||||
*/
|
||||
double red = (1 - percentage) * 255;
|
||||
|
||||
// Get all references to the function
|
||||
ReferenceIterator referenceIterator = currentProgram.getReferenceManager()
|
||||
.getReferencesTo(function.getEntryPoint());
|
||||
|
||||
// Iterate over all references
|
||||
for (Reference ref : referenceIterator) {
|
||||
// Create the colour
|
||||
Color color = new Color((int) red, 0, 0);
|
||||
// Set the background colour in the listing for the reference
|
||||
setBackgroundColor(ref.getFromAddress(), ref.getFromAddress(), color);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Takes a list of lists and returns a list of lists where none of the lists is
|
||||
* empty
|
||||
*
|
||||
* @param list the list to clean
|
||||
* @return the same list but without any empty lists within
|
||||
*/
|
||||
private List<List<Function>> cleanList(List<List<Function>> list) {
|
||||
List<List<Function>> output = new ArrayList<>();
|
||||
|
||||
for (List<Function> item : list) {
|
||||
if (item.isEmpty()) {
|
||||
continue;
|
||||
}
|
||||
output.add(item);
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
// Code below is taken from the Ghidra source code
|
||||
|
||||
public void setBackgroundColor(Address min, Address max, Color c) {
|
||||
IntRangeMap map = getColorRangeMap(true);
|
||||
if (map != null) {
|
||||
map.setValue(min, max, c.getRGB());
|
||||
}
|
||||
}
|
||||
|
||||
private IntRangeMap getColorRangeMap(boolean create) {
|
||||
IntRangeMap map = currentProgram.getIntRangeMap(PropertyBasedBackgroundColorModel.COLOR_PROPERTY_NAME);
|
||||
if (map == null && create) {
|
||||
try {
|
||||
map = currentProgram.createIntRangeMap(PropertyBasedBackgroundColorModel.COLOR_PROPERTY_NAME);
|
||||
} catch (DuplicateNameException e) {
|
||||
// can't happen since we just checked for it!
|
||||
}
|
||||
}
|
||||
return map;
|
||||
}
|
||||
|
||||
private int getMaxLevel(Map<CodeBlockVertex, Integer> levelMap) {
|
||||
int maxLevel = -1;
|
||||
for (Integer level : levelMap.values()) {
|
||||
if (level > maxLevel) {
|
||||
maxLevel = level;
|
||||
}
|
||||
}
|
||||
return maxLevel;
|
||||
}
|
||||
|
||||
private Function getFunctionFromCodeBlockVertex(CodeBlockVertex vertex) {
|
||||
Address startAddress = vertex.getCodeBlock().getFirstStartAddress();
|
||||
Function function = getFunctionAt(startAddress);
|
||||
return function;
|
||||
}
|
||||
|
||||
private List<List<Function>> createFunctionList(Map<CodeBlockVertex, Integer> levelMap) {
|
||||
List<List<Function>> levelList = new ArrayList<>();
|
||||
int maxLevel = getMaxLevel(levelMap);
|
||||
for (int i = 0; i <= maxLevel; i++) {
|
||||
levelList.add(new ArrayList<Function>());
|
||||
}
|
||||
for (CodeBlockVertex vertex : levelMap.keySet()) {
|
||||
int reverseLevel = maxLevel - levelMap.get(vertex);
|
||||
Function function = getFunctionFromCodeBlockVertex(vertex);
|
||||
if (function != null) {
|
||||
levelList.get(reverseLevel).add(function);
|
||||
}
|
||||
}
|
||||
return levelList;
|
||||
}
|
||||
|
||||
protected GDirectedGraph<CodeBlockVertex, CodeBlockEdge> createCallGraph() throws CancelledException {
|
||||
|
||||
Map<CodeBlock, CodeBlockVertex> instanceMap = new HashMap<>();
|
||||
GDirectedGraph<CodeBlockVertex, CodeBlockEdge> graph = GraphFactory.createDirectedGraph();
|
||||
|
||||
CodeBlockIterator codeBlocks = new BasicBlockModel(currentProgram, true).getCodeBlocks(monitor);
|
||||
while (codeBlocks.hasNext()) {
|
||||
CodeBlock block = codeBlocks.next();
|
||||
|
||||
CodeBlockVertex fromVertex = instanceMap.get(block);
|
||||
if (fromVertex == null) {
|
||||
fromVertex = new CodeBlockVertex(block);
|
||||
instanceMap.put(block, fromVertex);
|
||||
graph.addVertex(fromVertex);
|
||||
}
|
||||
|
||||
// destinations section
|
||||
addEdgesForDestinations(graph, fromVertex, block, instanceMap);
|
||||
}
|
||||
return graph;
|
||||
}
|
||||
|
||||
private void addEdgesForDestinations(GDirectedGraph<CodeBlockVertex, CodeBlockEdge> graph,
|
||||
CodeBlockVertex fromVertex, CodeBlock sourceBlock, Map<CodeBlock, CodeBlockVertex> instanceMap)
|
||||
throws CancelledException {
|
||||
|
||||
CodeBlockReferenceIterator iterator = sourceBlock.getDestinations(monitor);
|
||||
while (iterator.hasNext()) {
|
||||
monitor.checkCancelled();
|
||||
|
||||
CodeBlockReference destination = iterator.next();
|
||||
CodeBlock targetBlock = getDestinationBlock(destination);
|
||||
if (targetBlock == null) {
|
||||
continue; // no block found
|
||||
}
|
||||
|
||||
CodeBlockVertex targetVertex = instanceMap.get(targetBlock);
|
||||
if (targetVertex == null) {
|
||||
targetVertex = new CodeBlockVertex(targetBlock);
|
||||
instanceMap.put(targetBlock, targetVertex);
|
||||
}
|
||||
|
||||
graph.addVertex(targetVertex);
|
||||
graph.addEdge(new CodeBlockEdge(fromVertex, targetVertex));
|
||||
}
|
||||
}
|
||||
|
||||
private CodeBlock getDestinationBlock(CodeBlockReference destination) throws CancelledException {
|
||||
|
||||
Address targetAddress = destination.getDestinationAddress();
|
||||
CodeBlock targetBlock = new BasicBlockModel(currentProgram, true).getFirstCodeBlockContaining(targetAddress,
|
||||
monitor);
|
||||
if (targetBlock == null) {
|
||||
return null; // no code found for call; external?
|
||||
}
|
||||
|
||||
return targetBlock;
|
||||
}
|
||||
}
|
||||
+764
@@ -0,0 +1,764 @@
|
||||
//Gets all functions within the program and allows one to see the complexity depth of each function. The graph related code has been taken from and inspired by Ghidra's base: https://github.com/NationalSecurityAgency/ghidra/blob/master/Ghidra/Features/Base/src/main/java/ghidra/app/cmd/module/ComplexityDepthModularizationCmd.java#L43
|
||||
//@author Max 'Libra' Kersten for Trellix
|
||||
//@category
|
||||
//@keybinding
|
||||
//@menupath
|
||||
//@toolbar
|
||||
|
||||
import java.io.BufferedReader;
|
||||
import java.io.IOException;
|
||||
import java.io.InputStreamReader;
|
||||
import java.io.OutputStream;
|
||||
import java.net.HttpURLConnection;
|
||||
import java.net.URI;
|
||||
import java.net.URL;
|
||||
import java.util.ArrayList;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
|
||||
import com.google.gson.Gson;
|
||||
import com.google.gson.JsonSyntaxException;
|
||||
|
||||
import ghidra.app.decompiler.DecompInterface;
|
||||
import ghidra.app.decompiler.DecompileResults;
|
||||
import ghidra.app.decompiler.DecompiledFunction;
|
||||
import ghidra.app.plugin.core.analysis.AutoAnalysisManager;
|
||||
import ghidra.app.script.GhidraScript;
|
||||
import ghidra.graph.GDirectedGraph;
|
||||
import ghidra.graph.GraphAlgorithms;
|
||||
import ghidra.graph.GraphFactory;
|
||||
import ghidra.program.model.address.Address;
|
||||
import ghidra.program.model.block.BasicBlockModel;
|
||||
import ghidra.program.model.block.CodeBlock;
|
||||
import ghidra.program.model.block.CodeBlockIterator;
|
||||
import ghidra.program.model.block.CodeBlockReference;
|
||||
import ghidra.program.model.block.CodeBlockReferenceIterator;
|
||||
import ghidra.program.model.block.graph.CodeBlockEdge;
|
||||
import ghidra.program.model.block.graph.CodeBlockVertex;
|
||||
import ghidra.program.model.listing.Function;
|
||||
import ghidra.program.model.listing.Variable;
|
||||
import ghidra.program.model.symbol.SourceType;
|
||||
import ghidra.util.SystemUtilities;
|
||||
import ghidra.util.exception.CancelledException;
|
||||
|
||||
public class GhidrAI extends GhidraScript {
|
||||
|
||||
/*
|
||||
* The number of seconds the decompiler will run before timing out
|
||||
*/
|
||||
private final int DECOMPILER_TIMEOUT = 600;
|
||||
|
||||
/*
|
||||
* The number of threads via which the AI LLM provider is contacted. This
|
||||
* decreases the wait time on the LLM as multiple functions are handled at the
|
||||
* same time. If your LLM cannot handle multiple connections, adjust this value.
|
||||
*/
|
||||
private final int THREAD_COUNT = 10;
|
||||
|
||||
/*
|
||||
* The number of milliseconds to wait before the HTTP request times out
|
||||
*/
|
||||
private final int POST_REQUEST_TIMEOUT = 120_000;
|
||||
|
||||
/*
|
||||
* The HTTP URL to the LLM API
|
||||
*/
|
||||
private String API_URL;
|
||||
|
||||
/*
|
||||
* True if functions should be renamed, false if not. Only functions starting
|
||||
* with FUN_ are renamed if this is true.
|
||||
*/
|
||||
private boolean RENAME_FUNCTION;
|
||||
|
||||
/*
|
||||
* True if variables within functions should be renamed
|
||||
*/
|
||||
private boolean RENAME_VARIABLES;
|
||||
|
||||
/*
|
||||
* The Google JSON handling library which is included in Ghidra, used to convert
|
||||
* the JSON response from the LLM proxy
|
||||
*/
|
||||
private Gson gson;
|
||||
|
||||
/*
|
||||
* Ghidra's automatic analysis manager, used to check if one or more analyzers
|
||||
* are still running
|
||||
*/
|
||||
private AutoAnalysisManager autoAnalysisManager;
|
||||
|
||||
@Override
|
||||
protected void run() throws Exception {
|
||||
/*
|
||||
* Required input:
|
||||
*
|
||||
* - AI API URL
|
||||
*
|
||||
* - Rename function based on suggestion: yes/no
|
||||
*
|
||||
* - Rename variables based on suggestions: yes/no
|
||||
*/
|
||||
|
||||
// Get all provided values
|
||||
API_URL = askString("AI API URL", "The URL where the AI's API is accessible at");
|
||||
RENAME_FUNCTION = askYesNo("", "Rename the function based on the AI's suggestion");
|
||||
RENAME_VARIABLES = askYesNo("", "Rename variables based on the AI's suggestions");
|
||||
|
||||
// Create a string to provide the feedback to the user
|
||||
String inputFeedback = "Received input:\n";
|
||||
inputFeedback += "\tAI URL: " + API_URL + "\n";
|
||||
inputFeedback += "\tRename functions: " + RENAME_FUNCTION + "\n";
|
||||
inputFeedback += "\tRename variables: " + RENAME_VARIABLES;
|
||||
|
||||
// Log the input feedback
|
||||
println(inputFeedback);
|
||||
|
||||
// Now that all checks are done, the variables in the class are initialised
|
||||
gson = new Gson();
|
||||
autoAnalysisManager = AutoAnalysisManager.getAnalysisManager(currentProgram);
|
||||
DecompInterface decompiler = new DecompInterface();
|
||||
decompiler.openProgram(currentProgram);
|
||||
|
||||
/*
|
||||
* Notify the user of the current activity. Generating the complexity graph
|
||||
* takes a bit, but each step is relatively small. As such, it is summarised as
|
||||
* a single-step task to the user
|
||||
*/
|
||||
monitor.initialize(1, "Generating the complexity graph for the current program");
|
||||
|
||||
// Generate the complexity graph
|
||||
GDirectedGraph<CodeBlockVertex, CodeBlockEdge> callGraph = createCallGraph();
|
||||
Map<CodeBlockVertex, Integer> complexityDepth = GraphAlgorithms.getComplexityDepth(callGraph);
|
||||
List<List<Function>> partition = createFunctionList(complexityDepth);
|
||||
|
||||
// Remove the empty levels from the list
|
||||
partition = cleanAndReverseList(partition);
|
||||
|
||||
int currentFunctionCount = 0;
|
||||
int totalFunctionCount = 0;
|
||||
for (List<Function> list : partition) {
|
||||
totalFunctionCount += list.size();
|
||||
}
|
||||
|
||||
// Increment the monitor, marking this task as complete
|
||||
monitor.increment();
|
||||
|
||||
/*
|
||||
* To iterate over the list from start to end (rather than end to start), you
|
||||
* need to reverse the list. This can be done using
|
||||
* Collections.reverse(partition).
|
||||
*/
|
||||
|
||||
/*
|
||||
* Declare the executor service variable. It is instantiated per level within
|
||||
* the partition.
|
||||
*/
|
||||
ExecutorService executor;
|
||||
|
||||
/*
|
||||
* Iterate backwards over the list, handling the functions with the least
|
||||
* complexity depth first
|
||||
*/
|
||||
for (int i = 0; i < partition.size(); i++) {
|
||||
// Get the functions for the complexity level
|
||||
List<Function> list = partition.get(i);
|
||||
|
||||
if (list.isEmpty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
/*
|
||||
* Create the monitor message, taking the potential plurality of the function
|
||||
* count into account
|
||||
*/
|
||||
String listSize = "Iterating over " + list.size() + " function";
|
||||
if (list.size() > 1) {
|
||||
listSize += "s";
|
||||
}
|
||||
String monitorMessage = listSize + " in level " + i + "/" + partition.size();
|
||||
|
||||
/*
|
||||
* Initialise the monitor with the message. Each thread will call back to the
|
||||
* thread safe increment method
|
||||
*/
|
||||
monitor.initialize(list.size(), monitorMessage);
|
||||
|
||||
/*
|
||||
* The number of threads is equal to the globally set maximum number of thread
|
||||
* counts unless the size of the list is less than that.
|
||||
*/
|
||||
if (list.size() > THREAD_COUNT) {
|
||||
executor = Executors.newFixedThreadPool(THREAD_COUNT);
|
||||
} else {
|
||||
executor = Executors.newFixedThreadPool(list.size());
|
||||
}
|
||||
|
||||
// Iterate over each function within this complexity level
|
||||
for (Function function : list) {
|
||||
// Exclude thunk and external functions
|
||||
if (function.isThunk() || function.isExternal()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
/*
|
||||
* Set the current value for the current function, so we can keep the analyst
|
||||
* updated based on the number of functions (as only the current and total
|
||||
* number of levels in the graph is shown via the monitor message
|
||||
*/
|
||||
currentFunctionCount++;
|
||||
/*
|
||||
* Create a worker to schedule in the thread pool. The worker will connect with
|
||||
* the LLM and use the response to modify the given function based on the
|
||||
* predefined settings with regards to renaming.
|
||||
*/
|
||||
FunctionWorker worker = new FunctionWorker(decompiler, function, totalFunctionCount,
|
||||
currentFunctionCount);
|
||||
// Schedule the worker in the executor
|
||||
executor.execute(worker);
|
||||
}
|
||||
|
||||
// Shut the executor down, meaning no new workers can be added
|
||||
executor.shutdown();
|
||||
|
||||
// Wait until all workers have completed their work
|
||||
while (executor.isTerminated() == false) {
|
||||
// If the user cancels the script, cancel the execution
|
||||
if (monitor.isCancelled()) {
|
||||
throw new CancelledException();
|
||||
}
|
||||
/*
|
||||
* If the workers are running and the script isn't cancelled by the user, sleep
|
||||
* for one second to avoid using CPU cycles and check again after the sleep
|
||||
* finishes.
|
||||
*/
|
||||
Thread.sleep(1000);
|
||||
}
|
||||
/*
|
||||
* Before moving on to the next level, we wait for any ongoing automatic
|
||||
* analysis to complete
|
||||
*/
|
||||
waitUntilAutoAnalysisCompletes();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* A function to remove empty list entries from the list of lists. The functions
|
||||
* within each list within the given list are also iterated over, excluding
|
||||
* external and thunk functions, as these will not be handled afterwards
|
||||
*
|
||||
* @param list the list which contains the levels, where each level is a list of
|
||||
* functions
|
||||
* @return a list with the levels, where each level contains 1 or more
|
||||
* functions, and all included functions are non-thunk and non-external
|
||||
*/
|
||||
private List<List<Function>> cleanAndReverseList(List<List<Function>> list) {
|
||||
// The output variable
|
||||
List<List<Function>> output = new ArrayList<>();
|
||||
|
||||
// Iterate backwards over the list
|
||||
for (int i = list.size(); i-- > 0;) {
|
||||
// Get the current level
|
||||
List<Function> level = list.get(i);
|
||||
// Skip empty levels
|
||||
if (level.isEmpty()) {
|
||||
continue;
|
||||
}
|
||||
// Create a new list of functions
|
||||
List<Function> functions = new ArrayList<>();
|
||||
// Iterate over the functions within the level
|
||||
for (Function f : level) {
|
||||
// Exclude external and thunk functions
|
||||
if (f.isExternal() || f.isThunk()) {
|
||||
continue;
|
||||
}
|
||||
// Add non-thunk and non-external functions to the list
|
||||
functions.add(f);
|
||||
}
|
||||
// Add the filtered level to the output variable
|
||||
output.add(functions);
|
||||
}
|
||||
|
||||
/*
|
||||
* Return the list with the levels, in the original order with the applied
|
||||
* filters
|
||||
*/
|
||||
return output;
|
||||
}
|
||||
|
||||
/**
|
||||
* Prints the given message via the
|
||||
* {@link ghidra.app.script.GhidraScript#println(String)} method, depending on
|
||||
* the running mode and boolean.
|
||||
*
|
||||
* @param message the message to print
|
||||
* @param printHeadless true if the message should be printed when executing
|
||||
* headless, false if not
|
||||
*/
|
||||
private void log(String message, boolean printHeadless) {
|
||||
if (SystemUtilities.isInHeadlessMode()) {
|
||||
if (printHeadless == false) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
println(message);
|
||||
}
|
||||
|
||||
/**
|
||||
* This method only returns once the auto analysis manager is done analyzing. If
|
||||
* the user cancels, a cancelled exception is thrown.
|
||||
*
|
||||
* @throws InterruptedException
|
||||
* @throws CancelledException
|
||||
*/
|
||||
private void waitUntilAutoAnalysisCompletes() throws InterruptedException, CancelledException {
|
||||
while (autoAnalysisManager.isAnalyzing()) {
|
||||
if (monitor.isCancelled()) {
|
||||
throw new CancelledException();
|
||||
}
|
||||
Thread.sleep(1000); // Sleep to avoid consuming extra CPU cycles in-between checks
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* A thread safe method to increment the monitor
|
||||
*/
|
||||
private synchronized void incrementMonitorThreadSafe() {
|
||||
// Does not check for cancellation
|
||||
monitor.incrementProgress();
|
||||
}
|
||||
|
||||
/**
|
||||
* A thread safe method to verify if a method's name is unique.
|
||||
*
|
||||
* @param functionName The function name to verify
|
||||
* @return true if the name is unique, false if not
|
||||
*/
|
||||
private synchronized boolean isUniqueFunctionName(String functionName) {
|
||||
for (Function function : currentProgram.getFunctionManager().getFunctionsNoStubs(currentProgram.getMinAddress(),
|
||||
true)) {
|
||||
if (function.getName().equalsIgnoreCase(functionName)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a unique function name. If the name is not unique, an underscore is
|
||||
* appended
|
||||
*
|
||||
* @param functionName the name to check
|
||||
* @return the unique name
|
||||
*/
|
||||
private synchronized String getUniqueFunctionName(String functionName) {
|
||||
if (isUniqueFunctionName(functionName)) {
|
||||
return functionName;
|
||||
}
|
||||
|
||||
functionName += "_";
|
||||
return getUniqueFunctionName(functionName);
|
||||
}
|
||||
|
||||
/**
|
||||
* The class where the threaded execution to handle a given function is handled
|
||||
*/
|
||||
private class FunctionWorker implements Runnable {
|
||||
|
||||
private DecompInterface decompiler;
|
||||
private Function function;
|
||||
int totalFunctionCount;
|
||||
int currentFunctionCount;
|
||||
|
||||
public FunctionWorker(DecompInterface decompiler, Function function, int totalFunctionCount,
|
||||
int currentFunctionCount) {
|
||||
this.decompiler = decompiler;
|
||||
this.function = function;
|
||||
this.totalFunctionCount = totalFunctionCount;
|
||||
this.currentFunctionCount = currentFunctionCount;
|
||||
}
|
||||
|
||||
/**
|
||||
* The run method of the thread, not to be confused with the script's main run
|
||||
* method!
|
||||
*/
|
||||
@Override
|
||||
public void run() {
|
||||
// Decompile the given function
|
||||
DecompileResults results = decompiler.decompileFunction(function, DECOMPILER_TIMEOUT, monitor);
|
||||
// Get the decompiled function
|
||||
DecompiledFunction dFunction = results.getDecompiledFunction();
|
||||
// Get the pseudo-C representation of the decompiled function
|
||||
String code = dFunction.getC();
|
||||
// Get the plate comment, if present
|
||||
String tempComment = getPlateComment(function.getEntryPoint());
|
||||
// If the comment is non-null, it exists
|
||||
if (tempComment != null) {
|
||||
/*
|
||||
* Remove the comment from the pseudo-C representation to avoid overloading the
|
||||
* LLM's context window
|
||||
*/
|
||||
code = code.replace(tempComment, "");
|
||||
}
|
||||
|
||||
// Get the function's variables
|
||||
Variable[] variables = function.getAllVariables();
|
||||
|
||||
// Contact the LLM and store the JSON-based response
|
||||
JsonResponse response = contactLLM(code);
|
||||
// If the response is null, an error occurred
|
||||
if (response == null) {
|
||||
// Print an error if there is an error
|
||||
printerr(
|
||||
"The LLM's return value cannot be parsed as the format is invalid or the LLM did not respond at all");
|
||||
// Return from the thread
|
||||
return;
|
||||
}
|
||||
|
||||
// Declare and initialize the comment
|
||||
String comment = "";
|
||||
|
||||
// Store the old function name
|
||||
String oldFunctionName = function.getName();
|
||||
// Get the new function name from the LLM response
|
||||
String newFunctionName = response.getFunctionName();
|
||||
// If the new name is present and is not an empty or whitespace-only string
|
||||
if (newFunctionName != null && newFunctionName.isBlank() == false) {
|
||||
// If functions should be renamed
|
||||
if (RENAME_FUNCTION) {
|
||||
// Only rename functions with default names
|
||||
if (function.getName().toLowerCase().startsWith("fun_")) {
|
||||
try {
|
||||
// Set the new name
|
||||
function.setName(getUniqueFunctionName(newFunctionName), SourceType.IMPORTED);
|
||||
// Log the rename to non-headless instances
|
||||
log(oldFunctionName + " -> " + newFunctionName, false);
|
||||
// Add the name change to the comment
|
||||
comment += "Changed " + oldFunctionName + " into " + newFunctionName + "\n\n";
|
||||
} catch (Exception ex) {
|
||||
printerr("Function renaming failed:\n" + ex.toString());
|
||||
}
|
||||
}
|
||||
}
|
||||
/*
|
||||
* If the function should not be renamed, or if the function name does not start
|
||||
* with fun_, the suggested name is still added to the comment
|
||||
*/
|
||||
if (RENAME_FUNCTION == false || function.getName().toLowerCase().startsWith("fun_") == false) {
|
||||
comment += "AI suggested function name: " + newFunctionName + "\n\n";
|
||||
}
|
||||
}
|
||||
|
||||
// If one or more variables are present in the function and LLM response
|
||||
if (variables.length > 0 && response.getVariableNames() != null) {
|
||||
// Set the next part of the comment
|
||||
comment += "Variables:\n";
|
||||
// Iterate over all variables
|
||||
for (Variable variable : variables) {
|
||||
// Get the old name
|
||||
String oldVariableName = variable.getName();
|
||||
// Get the new name
|
||||
String newVariableName = response.getVariableNames().get(oldVariableName);
|
||||
// If the new name is non-null and is not an empty or whitespace-only string
|
||||
if (newVariableName != null && newVariableName.isBlank() == false) {
|
||||
// If variables should be renamed
|
||||
if (RENAME_VARIABLES) {
|
||||
try {
|
||||
// Set the name
|
||||
variable.setName(newVariableName, SourceType.IMPORTED);
|
||||
} catch (Exception ex) {
|
||||
printerr("Variable renaming failed:\n" + ex.toString());
|
||||
}
|
||||
}
|
||||
// Add the renaming to the comment
|
||||
comment += "\t\tAI variable name: \"" + oldVariableName + "\" -> \"" + newVariableName + "\"\n";
|
||||
}
|
||||
}
|
||||
// Log the variables in Ghidra, when running non-headless
|
||||
log("Variables for " + function.getName() + " renamed", false);
|
||||
comment += "\n";
|
||||
}
|
||||
|
||||
// Get the summary from the LLM
|
||||
String summary = response.getSummary();
|
||||
// If the summary is non-null and is not an empty or whitespace-only string
|
||||
if (summary != null && summary.isBlank() == false) {
|
||||
// Add the summary to the comment
|
||||
comment += "AI function summary: " + summary;
|
||||
}
|
||||
|
||||
// If the comment is not whitespace-only or empty
|
||||
if (comment.isBlank() == false) {
|
||||
// Set the comment to the function
|
||||
setComment(function, comment, false);
|
||||
// Log the summary addition for non-headless instances
|
||||
log("Summary for " + function.getName() + " added", false);
|
||||
}
|
||||
|
||||
/*
|
||||
* This message is always logged (when running headless and non-headless) to
|
||||
* give an indication of the progress
|
||||
*/
|
||||
log("Completed " + function.getName() + " (" + currentFunctionCount + "/" + totalFunctionCount + ")", true);
|
||||
// Increment the monitor to update the progress in a thread safe manner
|
||||
incrementMonitorThreadSafe();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Send the HTTP POST request to LLM
|
||||
*
|
||||
* @param body the body of comma separated strings to check
|
||||
* @return the JSON response from the server
|
||||
* @throws IOException
|
||||
*/
|
||||
private String sendPostRequest(Map<String, String> headers, String body) throws IOException {
|
||||
URL url = URI.create(API_URL).toURL();
|
||||
HttpURLConnection connection = (HttpURLConnection) url.openConnection();
|
||||
connection.setRequestMethod("POST");
|
||||
|
||||
for (Map.Entry<String, String> entry : headers.entrySet()) {
|
||||
connection.setRequestProperty(entry.getKey(), entry.getValue());
|
||||
}
|
||||
|
||||
connection.setConnectTimeout(POST_REQUEST_TIMEOUT);
|
||||
connection.setReadTimeout(POST_REQUEST_TIMEOUT);
|
||||
|
||||
connection.setDoOutput(true);
|
||||
try (OutputStream outputStream = connection.getOutputStream()) {
|
||||
byte[] rawBody = body.getBytes("utf-8");
|
||||
outputStream.write(rawBody, 0, rawBody.length);
|
||||
}
|
||||
|
||||
// TODO use response code to potentially throw an exception
|
||||
int responseCode = connection.getResponseCode();
|
||||
|
||||
try (BufferedReader bufferedReader = new BufferedReader(
|
||||
new InputStreamReader(connection.getInputStream(), "utf-8"))) {
|
||||
StringBuilder response = new StringBuilder();
|
||||
String line = null;
|
||||
while ((line = bufferedReader.readLine()) != null) {
|
||||
response.append(line.trim());
|
||||
}
|
||||
return response.toString();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Prompt the LLM by providing the pseudo-C code as the argument
|
||||
*
|
||||
* @param code the function's pseudo-C code
|
||||
* @return the LLM's JSON response
|
||||
* @throws IOException
|
||||
*/
|
||||
public String promptLlmApi(String code) throws IOException {
|
||||
Map<String, String> headers = new HashMap<>();
|
||||
headers.put("Content-Type", "application/json");
|
||||
headers.put("Accept", "application/json");
|
||||
|
||||
String prompt = "For the following decompiled code from Ghidra, suggest a new name for the function, summarise the function, and suggest new names for each of the variables in the function. State the old name of the variable and the new name. The response should be a JSON object with \"functionName\" as the key for the suggested function name, the key \"summary\" should contain the function's summary, and a nested JSON object named \"variableNames\" with the old variable names as keys, where the value of each key is the newly suggested name. The response should only the the requested JSON object, nothing else.\n\n"
|
||||
+ code;
|
||||
|
||||
return sendPostRequest(headers, prompt);
|
||||
}
|
||||
|
||||
/**
|
||||
* Contact the LLM with the pseudo-C code and prompt
|
||||
*
|
||||
* @param code the function's pseudo-C code
|
||||
* @return a Java object in which the LLM's JSON response is already parsed
|
||||
*/
|
||||
private JsonResponse contactLLM(String code) {
|
||||
try {
|
||||
String response = promptLlmApi(code);
|
||||
|
||||
if (response == null) {
|
||||
// return null;
|
||||
}
|
||||
JsonResponse jsonResponse = gson.fromJson(response, JsonResponse.class);
|
||||
return jsonResponse;
|
||||
} catch (IOException ex) {
|
||||
printerr(ex.toString());
|
||||
return null;
|
||||
} catch (JsonSyntaxException ex) {
|
||||
printerr(ex.toString());
|
||||
return null;
|
||||
} catch (Exception ex) {
|
||||
printerr(ex.toString());
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Helper function to set a comment at a given function, with the indication if
|
||||
* this comment should be at the top or bottom of any already existing function
|
||||
* comment at this function
|
||||
*
|
||||
* @param function the function to set a comment at
|
||||
* @param comment the comment to set at the given function
|
||||
* @param pushToTop true if the function's existing comment (if any) should be
|
||||
* put below this comment, false if not
|
||||
*/
|
||||
private void setComment(Function function, String comment, boolean pushToTop) {
|
||||
// Get the old comment, which is null if no comment is present
|
||||
String oldComment = function.getComment();
|
||||
// Check for the existence of the old comment
|
||||
if (oldComment != null) {
|
||||
// If it is present, trim the comment to avoid redundant whitespace
|
||||
oldComment = oldComment.trim();
|
||||
// If the string, post trimming, is not empty nor blank (which also checks if it
|
||||
// is empty)
|
||||
if (oldComment.isEmpty() == false && oldComment.isBlank() == false) {
|
||||
// If the comment should be at the top, place it at the top
|
||||
if (pushToTop) {
|
||||
comment += "\n\n" + oldComment;
|
||||
} else { // Else put it at the bottom
|
||||
comment = oldComment += "\n\n" + comment;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the comment, which contains the prior comment at the correct placement
|
||||
* with regards to the given boolean if it existed
|
||||
*/
|
||||
function.setComment(comment);
|
||||
}
|
||||
|
||||
/**
|
||||
* The class used to store the HTTP LLM API response in
|
||||
*/
|
||||
private class JsonResponse {
|
||||
private String functionName;
|
||||
private String summary;
|
||||
private Map<String, String> variableNames;
|
||||
|
||||
@SuppressWarnings("unused")
|
||||
public JsonResponse() {
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets the AI generated function name
|
||||
*
|
||||
* @return the AI generated function name
|
||||
*/
|
||||
public String getFunctionName() {
|
||||
return functionName;
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets the AI generated summary
|
||||
*
|
||||
* @return the AI generated summary
|
||||
*/
|
||||
public String getSummary() {
|
||||
return summary;
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets the mapping of the variable names (as keys) and the AI generated
|
||||
* variable names (as a value, one for each key)
|
||||
*
|
||||
* @return gets the AI generated variable names for the current variable names
|
||||
*/
|
||||
public Map<String, String> getVariableNames() {
|
||||
return variableNames;
|
||||
}
|
||||
}
|
||||
|
||||
// From here onwards, the code is from the Ghidra source code
|
||||
|
||||
private int getMaxLevel(Map<CodeBlockVertex, Integer> levelMap) {
|
||||
int maxLevel = -1;
|
||||
for (Integer level : levelMap.values()) {
|
||||
if (level > maxLevel) {
|
||||
maxLevel = level;
|
||||
}
|
||||
}
|
||||
return maxLevel;
|
||||
}
|
||||
|
||||
private Function getFunctionFromCodeBlockVertex(CodeBlockVertex vertex) {
|
||||
Address startAddress = vertex.getCodeBlock().getFirstStartAddress();
|
||||
Function function = getFunctionAt(startAddress);
|
||||
return function;
|
||||
}
|
||||
|
||||
private List<List<Function>> createFunctionList(Map<CodeBlockVertex, Integer> levelMap) {
|
||||
List<List<Function>> levelList = new ArrayList<>();
|
||||
int maxLevel = getMaxLevel(levelMap);
|
||||
for (int i = 0; i <= maxLevel; i++) {
|
||||
levelList.add(new ArrayList<Function>());
|
||||
}
|
||||
for (CodeBlockVertex vertex : levelMap.keySet()) {
|
||||
int reverseLevel = maxLevel - levelMap.get(vertex);
|
||||
Function function = getFunctionFromCodeBlockVertex(vertex);
|
||||
if (function != null) {
|
||||
levelList.get(reverseLevel).add(function);
|
||||
}
|
||||
}
|
||||
return levelList;
|
||||
}
|
||||
|
||||
protected GDirectedGraph<CodeBlockVertex, CodeBlockEdge> createCallGraph() throws CancelledException {
|
||||
|
||||
Map<CodeBlock, CodeBlockVertex> instanceMap = new HashMap<>();
|
||||
GDirectedGraph<CodeBlockVertex, CodeBlockEdge> graph = GraphFactory.createDirectedGraph();
|
||||
CodeBlockIterator codeBlocks = new BasicBlockModel(currentProgram, true).getCodeBlocks(monitor);
|
||||
while (codeBlocks.hasNext()) {
|
||||
CodeBlock block = codeBlocks.next();
|
||||
|
||||
CodeBlockVertex fromVertex = instanceMap.get(block);
|
||||
if (fromVertex == null) {
|
||||
fromVertex = new CodeBlockVertex(block);
|
||||
instanceMap.put(block, fromVertex);
|
||||
graph.addVertex(fromVertex);
|
||||
}
|
||||
|
||||
// destinations section
|
||||
addEdgesForDestinations(graph, fromVertex, block, instanceMap);
|
||||
}
|
||||
return graph;
|
||||
}
|
||||
|
||||
private void addEdgesForDestinations(GDirectedGraph<CodeBlockVertex, CodeBlockEdge> graph,
|
||||
CodeBlockVertex fromVertex, CodeBlock sourceBlock, Map<CodeBlock, CodeBlockVertex> instanceMap)
|
||||
throws CancelledException {
|
||||
|
||||
CodeBlockReferenceIterator iterator = sourceBlock.getDestinations(monitor);
|
||||
while (iterator.hasNext()) {
|
||||
monitor.checkCancelled();
|
||||
|
||||
CodeBlockReference destination = iterator.next();
|
||||
CodeBlock targetBlock = getDestinationBlock(destination);
|
||||
if (targetBlock == null) {
|
||||
continue; // no block found
|
||||
}
|
||||
|
||||
CodeBlockVertex targetVertex = instanceMap.get(targetBlock);
|
||||
if (targetVertex == null) {
|
||||
targetVertex = new CodeBlockVertex(targetBlock);
|
||||
instanceMap.put(targetBlock, targetVertex);
|
||||
}
|
||||
|
||||
graph.addVertex(targetVertex);
|
||||
graph.addEdge(new CodeBlockEdge(fromVertex, targetVertex));
|
||||
}
|
||||
}
|
||||
|
||||
private CodeBlock getDestinationBlock(CodeBlockReference destination) throws CancelledException {
|
||||
|
||||
Address targetAddress = destination.getDestinationAddress();
|
||||
CodeBlock targetBlock = new BasicBlockModel(currentProgram, true).getFirstCodeBlockContaining(targetAddress,
|
||||
monitor);
|
||||
if (targetBlock == null) {
|
||||
return null; // no code found for call; external?
|
||||
}
|
||||
|
||||
return targetBlock;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user