reman3/scripts/re3lib/PCallTracer.java

package re3lib;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;

import ghidra.app.decompiler.DecompileResults;
import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.lang.Register;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.FunctionManager;
import ghidra.program.model.listing.Instruction;
import ghidra.program.model.listing.InstructionIterator;
import ghidra.program.model.listing.Program;
import ghidra.program.model.pcode.HighFunction;
import ghidra.program.model.pcode.PcodeOp;
import ghidra.program.model.pcode.PcodeOpAST;
import ghidra.program.model.pcode.Varnode;
import ghidra.program.model.symbol.Reference;

public class PCallTracer {
	public class QueueItem {
		Function function;
		int depth;

		QueueItem(Function function, int depth) {
			this.function = function;
			this.depth = depth;
		}
	}

	public List<Function> out = new ArrayList<>();
	public boolean trace = false;
	List<QueueItem> queue = new ArrayList<>();
	HashSet<Address> visited = new HashSet<>();

	// DecompileCache decomp;
	GhidraScript script;
	Program program;

	public PCallTracer() {
		this.script = RecompileConfig.INSTANCE.script;
		this.program = this.script.getCurrentProgram();
		// this.decomp = RecompileConfig.INSTANCE.decompCache;
	}

	public void setBlacklist(HashSet<Address> blacklist) {
		this.visited = new HashSet<>(blacklist);
	}

	// public List<Address> getReferenceFromAddresses(Address address) {
	// Reference[] referencesFrom = script.getReferencesFrom(address);

	// // get only the address references at the given address (ie no stack refs,
	// ...)
	// List<Address> refFromAddresses = new ArrayList<Address>();
	// for (Reference referenceFrom : referencesFrom) {
	// if (referenceFrom.isMemoryReference()) {
	// refFromAddresses.add(referenceFrom.getToAddress());
	// }
	// }

	// return refFromAddresses;
	// }

	public Function getReferencedFunction(Address address, boolean getThunkedFunction) {
		Reference[] referencesFrom = script.getReferencesFrom(address);

		if (referencesFrom.length == 0) {
			return null;
		}

		for (Reference referenceFrom : referencesFrom) {
			Address referencedAddress = referenceFrom.getToAddress();
			if (referencedAddress == null) {
				continue;
			}

			Function function = script.getFunctionAt(referencedAddress);

			if (function == null) {
				continue;
			}

			if (!getThunkedFunction) {
				return function;
			}

			if (function.isThunk()) {
				function = function.getThunkedFunction(true);
			}
			return function;
		}

		return null;
		// List<Address> referencesFrom = getReferenceFromAddresses(address);
		// if (referencesFrom.size() != 1) {
		// return null;
		// }

		// Address functionAddress = referencesFrom.get(0);

		// Register lowBitCodeMode = program.getRegister("LowBitCodeMode");
		// if (lowBitCodeMode != null) {
		// long longValue = functionAddress.getOffset();
		// longValue = longValue & ~0x1;
		// functionAddress = functionAddress.getNewAddress(longValue);
		// }

		// Function function = script.getFunctionAt(functionAddress);
		// if (function == null) {
		// // try to create function
		// function = script.createFunction(functionAddress, null);
		// if (function == null) {
		// return null;
		// }
		// }

		// // if function is a thunk, get the thunked function
		// if (function.isThunk()) {
		// Function thunkedFunction = function.getThunkedFunction(true);
		// function = thunkedFunction;
		// }

		// return function;
	}

	void visitInstructions(Function function, int depth) {
		Set<Function> calledFunctions = new HashSet<Function>();

		InstructionIterator instructions = script.getCurrentProgram()
				.getListing()
				.getInstructions(function.getBody(), true);
		while (instructions.hasNext()) {
			Instruction instruction = instructions.next();
			if (instruction.getFlowType().isCall()) {

				FunctionManager functionManager = script.getCurrentProgram().getFunctionManager();
				Function calledFunction = getReferencedFunction(instruction.getMinAddress(), true);
				if (calledFunction == null) {
					continue;
				}

				if (!visited.contains(calledFunction.getEntryPoint())) {
					queue.add(new QueueItem(calledFunction, depth + 1));
				}
			}
		}

		// Iterator<PcodeOpAST> opIter = highFunction.getPcodeOps();
		// while (opIter.hasNext()) {
		// PcodeOpAST op = opIter.next();
		// if (op.getOpcode() == PcodeOp.CALL) {
		// Varnode target = op.getInput(0);
		// if (target.isAddress()) {
		// Address callAddr = target.getAddress();
		// Function calledFunction = script.getFunctionAt(callAddr);
		// if (calledFunction == null) {
		// script.println("PCallTracer, called function not found: " + op.toString() + "
		// - "
		// + highFunction.getFunction().getName());
		// continue;
		// }
		// if (!visited.contains(calledFunction.getEntryPoint())) {
		// queue.add(new QueueItem(calledFunction, depth + 1));
		// }
		// }
		// }
		// }
	}

	void visit(Function function, int depth) {
		if (!visited.contains(function.getEntryPoint())) {
			visited.add(function.getEntryPoint());
			if (trace) {
				script.println("PCallTracer, visiting " + function.getName() + " (depth:" + depth + ")");
			}
			// DecompileResults decompRes = decomp.getOrInsert(function);
			visitInstructions(function, depth);
			out.add(function);
		}
	}

	public void traceCalls(Function inFunction) {
		queue.add(new QueueItem(inFunction, 0));
		while (queue.size() > 0) {
			QueueItem item = queue.remove(0);
			visit(item.function, item.depth);
		}
	}
}