Binary Ninja plugin that highlights instructions based on a csv file

load_cov.py

from binaryninja import PluginCommand, highlight, enums
from binaryninja.interaction import get_open_filename_input

from collections import namedtuple

FunctionData = namedtuple("FunctionData", ["function", "counts"])

g_highlights = dict()

def get_functions(view, address):
    funcs = view.get_functions_at(address)
    if len(funcs) == 0:
        address = view.get_previous_function_start_before(address)
        funcs = view.get_functions_at(address)

    if len(funcs) > 1:
        print("WARNING: more than one function at {addr:x}".format(
            addr=address
        ))

    if len(funcs) == 0:
        return None
    return funcs

def load_dict(filename):
    res = dict()
    f = open(filename, "r")
    for line in f:
        line = line.strip()
        if line == "":
            continue
        addr, count = line.split(",")
        addr = int(addr, 16)
        count = int(count)
        res[addr] = count
    f.close()
    return res

def populate_data(bv, d):
    functions = dict()
    for addr in d:
        fs = get_functions(bv, addr)
        if fs is None:
            print("NO FUNCTION FOR ", hex(addr))
            continue
        for f in fs:
            if f.name in functions:
                functions[f.name].counts.append(
                    (addr, d[addr])
                )
            else:
                functions[f.name] = FunctionData(
                    function=f,
                    counts=[
                        (addr, d[addr])
                    ]
                )
    return functions

def clear_cov(bv):
    if bv not in g_highlights:
        return

    for f, addr in g_highlights[bv]:
        f.set_auto_instr_highlight(addr, enums.HighlightStandardColor(0))
    g_highlights[bv].clear()

def display_cov(bv):
    file_path = get_open_filename_input("Filename: ")
    d = load_dict(file_path)

    if bv not in g_highlights:
        g_highlights[bv] = list()
    clear_cov(bv)
    functions = populate_data(bv, d)
    for f_name in functions:
        f_data = functions[f_name]
        f = f_data.function
        max_count = 0
        for _, count in f_data.counts:
            if count > max_count:
                max_count = count

        for addr, count in f_data.counts:
            gradient = int(80 + (255 - 80) * (count / max_count)) if max_count > 1 else 80
            color = highlight.HighlightColor(red=gradient, green=0, blue=0)
            # print("Setting highlight at address: ", hex(addr), "color: ", color, "gradient:", gradient)
            f.set_auto_instr_highlight(addr, color)
            g_highlights[bv].append((f, addr))

PluginCommand.register(
    "HighlightAddr\\Load CSV File",
    """
Highlight instructions contained in the CSV file: for each line, there must be a couple
"addr, num" where "num" indicates a score (greater the score, lighter the color in the highlight)
""",
    display_cov
)

PluginCommand.register(
    "HighlightAddr\\Clear",
    "clear highlight",
    clear_cov
)