Painful IDA coloring

script.py

import idc
import idaapi
import sark

base = idc.get_name_ea(0xbadec0de,"SepMandatoryIntegrityCheck")

arr_10 = [0x52, 0x63, 0x6f, 0x78, 0x85, 0x98, 0xa2, 0xad, 0xd2, 0xe2, 0xf3, 0x103, 0x10c, 0x119, 0x148, 0x167, 0x16e, 0x18d, 0x19b, 0x1a8, 0x1b6, 0x1c8, 0x1dc, 0x1e3, 0x1ed, 0x1f6, 0x20d, 0x21c, 0x3e2, 0x22e, 0x251, 0x31c, 0x26f, 0x340, 0x356, 0x35e, 0x36a, 0x38e, 0x39b, 0x295, 0x2cb, 0x2db]
executed_nodes_10 = [sark.CodeBlock(base+i).start_ea for i in arr_10]
arr_8 = [0x52, 0x63, 0x6f, 0x78, 0x85, 0x98, 0xa2, 0xad, 0xd2, 0xe2, 0xf3, 0x103, 0x10c, 0x119, 0x148, 0x167, 0x16e, 0x314, 0x22e, 0x251, 0x26f, 0x295, 0x2cb, 0x2db]
executed_nodes_8 = [sark.CodeBlock(base+i).start_ea for i in arr_8]

common = set(executed_nodes_10) & set(executed_nodes_8)
already_colored = set()


for i in arr_10:
    already_colored.add(i)
    sark.CodeBlock(base+i).color=0x0000ff

for i in arr_8:
    if i in already_colored:
        sark.CodeBlock(base+i).color=0x00ff00
    else:
        sark.CodeBlock(base+i).color=0xff0000

cfunc = idaapi.decompile(base)

def lex_citem_indexes(line):
    i = 0
    indexes = []
    line_length = len(line)
    while i < line_length:
        if line[i] == idaapi.COLOR_ON:
            i += 1
            if ord(line[i]) == idaapi.COLOR_ADDR:
                i += 1
                citem_index = int(line[i:i+idaapi.COLOR_ADDR_SIZE], 16)
                i += idaapi.COLOR_ADDR_SIZE
                indexes.append(citem_index)
                continue
        i += 1
    return indexes


def map_line2citem(decompilation_text):
    line2citem = {}
    for line_number in xrange(decompilation_text.size()):
        line_text = decompilation_text[line_number].line
        line2citem[line_number] = lex_citem_indexes(line_text)
    return line2citem

def map_line2node(cfunc, line2citem):
    line2node = {}
    treeitems = cfunc.treeitems
    function_address = cfunc.entry_ea
    for line_number, citem_indexes in line2citem.iteritems():
        nodes = set()
        for index in citem_indexes:
            try:
                item = treeitems[index]
                address = item.ea
            except IndexError as e:
                continue          
            if address == 0xffffffffffffffff:
                continue
            node = sark.CodeBlock(address).start_ea
            if not node:
                continue
            nodes.add(node)
        line2node[line_number] = nodes
    return line2node

def color(arr, color):
    decompilation_text = cfunc.get_pseudocode()

    line2citem = map_line2citem(decompilation_text)

    line2node = map_line2node(cfunc, line2citem)

    lines_painted = 0

    executed_nodes = set(arr)

    for line_number, line_nodes in line2node.iteritems():
        if line_nodes & executed_nodes:
            decompilation_text[line_number].bgcolor = color
            lines_painted += 1

    for line_number in xrange(0, cfunc.hdrlines):
        decompilation_text[line_number].bgcolor = color
        lines_painted += 1
    
    idaapi.refresh_idaview_anyway()


color(executed_nodes_10, 0x0000ff)
color(executed_nodes_8, 0xff00ff)
color(common, 0x00ff00)