Ghidra Script to Colorize all Basic Blocks for a Provided list of Addresses

color_bb.py

from java.awt import Color
from ghidra.util.task import ConsoleTaskMonitor
from ghidra.program.model.block import BasicBlockModel
from docking.options.editor import GhidraColorChooser

"""
Ghidra script to colorize all basic blocks identified by the input file.
Prompts for a file to use as input. This script will attempt to sanity check
that a basic block actually does start at each provided address.

Input should be a line-separated list of offsets from start of the program:

$ cat bb_list.txt
0x000001de02
0x000001de19
0x000001de1f
0x000001de27
0x000001de47
0x000001de4c
0x000001de53
0x000001de66
0x000001de6c
0x000001de7e
0x000001de98
0x000001deb4
"""

def colorize_bb(bb_address_list, color):
    bm = BasicBlockModel(currentProgram)
    monitor = ConsoleTaskMonitor()
    for addr in bb_address_list:
        addr_obj = toAddr(addr)
        if addr is None:
            print("Invalid address: {}", addr)
            continue
        bb = bm.getCodeBlockAt(addr_obj, monitor)
        if bb is None:
            print("No basic block at add ress: {:#018x}".format(addr))
            continue
        setBackgroundColor(bb, color)


def run():
    bb_filename = askFile("Select Basic Block Address List", "whatever")
    bb_filename = str(bb_filename)
    _color = Color(0xcc, 0xff, 0xcc)
    color = GhidraColorChooser.showDialog(None, "Choose a Color", _color)
    print(color)
    if color is None:
        color = _color

    bb_addrs = []
    with open(bb_filename, "r") as f:
        for line in f:
            try:
                addr = int(line.strip(), 16)
                bb_addrs.append(addr)
            except ValueError:
                continue
    colorize_bb(bb_addrs, color)


run()