//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;
	}
}