ida_basicblock_simple.py

• Linux path to script:
  • Not shipped on Linux (IDA script)
• Windows path to script:
  • amd64: C:\Program Files (x86)\Reven-PythonAPI-amd64\python-examples\ida\ida_basicblock_simple.py
  • x86: C:\Program Files (x86)\Reven-PythonAPI-x86\python-examples\ida\ida_basicblock_simple.py

This is an example of IDA script using the REVEN Python API. Must be run from IDA!

See the script's documentation below for more information:

"""
Purpose:
    This script is a proof-of-concept for simple program profiling: each basic
    block in a program will be colorized depending its execution frequency, one
    can also observe basic block frequencies in the output window of IDA.

Usage:
    Use IDA to load the binary, and give the following arguments before
    executing the script:
        host = your REVEN server name, and
        port = your REVEN's project port on this host.

Remark:
    For each basic block obtained from IDA's static analysis, the script counts
    its occurrence frequency in the trace given by REVEN. The following limits
    are known:
        - the name of the binary must be obtained from REVEN (that is usually the
          case if "dump process" is executed properly), and be case insentively
          identical with the one analyzed in IDA (i.e. we do not rename the binary)

        - the binary must be mapped at a unique virtual base address in the REVEN's
          project trace

        - very poor results for packed (or self-modifying) binaries (since the
          limit of IDA's static analysis)
"""

import idaapi
import idautils
import idc

import reven



def main(host, port):
    project = reven.Project(host, port)

    runtime_base_address = get_base_address(project)
    if runtime_base_address is not None:
        static_base_address = idaapi.get_imagebase()

        print 'base addresses:'
        print '  static: 0x{:x}'.format(idaapi.get_imagebase())
        print '  runtime: 0x{:x}'.format(runtime_base_address)

        offset = runtime_base_address - static_base_address

        static_bbs = collect_basic_blocks()
        binary_trace = get_binary_trace(project)
        dynamic_bbs = calculate_dynamic_basic_blocks(static_bbs, binary_trace, offset)

        bb_freq = calculate_dynamic_basic_block_frequency(dynamic_bbs, binary_trace, offset)
        colorize_dynamic_basic_blocks(dynamic_bbs, bb_freq)

        print 'colorizing basic blocks done.'

    else:
        print 'cannot find the binary {:s} in the REVEN\'s trace or it is mapped at different base addresses.'.format(os.path.basename(idc.GetInputFilePath()))

def get_base_address(reven_project):
    """
    Look for the current IDA's binary name in the REVEN's binary mapping information.
    If the binary is found, return the base address it is mapped at, else return None.
    """
    base_address = None
    binary_name = os.path.basename(idc.GetInputFilePath()).lower() # The path information is completely irrelevant,
                                                                   # so we will match against the binary name only
    loaded_binaries = reven_project.binaries()
    for bin_path in loaded_binaries:
        if binary_name == os.path.basename(bin_path).lower():
            bin_mappings = loaded_binaries[bin_path].mappings

            # Reven stores all the base addresses this binary is mapped at during the trace.
            # There could be more than one if more than one process uses this binary
            for address_space in bin_mappings.values():
                if base_address is None:
                    base_address = address_space.base_address
                elif base_address != address_space.base_address:
                    # More than one process uses the binary, and at different addresses!
                    # We don't have enough information to pick one.
                    return None

    return base_address


def get_binary_trace(reven_project):
    """
    Return a python generator that filters the trace on the current IDA's binary only.
    """
    # Select the main trace
    trace = reven_project.trace('Execution run')
    # Or in a better way:
    # trace = project.traces()[0]

    # Again, the binary's path is irrelevant, match against its name only.
    binary_name = os.path.basename(idc.GetInputFilePath())

    # See the documentation for more information, but note that this defaults to a "contains" matching algorithm.
    # It may therfore include more that the binary we want.
    reven_points = trace.search_point([reven.BinaryCriterion(pattern=binary_name,
                                                             case_sensitive=False)])
    return list(reven_points)


def calculate_dynamic_basic_block_frequency(dynamic_bbs, reven_trace, relocation_offset):
    bb_freq = { fst + relocation_offset: 0 for (fst, snd) in dynamic_bbs }

    for head_point in reven_trace:
        bb = head_point.basic_block

        for ins in bb:
            ins_addr = ins.address
            if ins_addr in bb_freq:
                bb_freq[ins_addr] += 1

    return { addr - relocation_offset: bb_freq[addr] for addr in bb_freq }


def calculate_frequency(ida_bbs, reven_trace, relocation_offset):
    bb_freq = {bb.startEA + relocation_offset: 0 for bb in ida_bbs}

    # The generator we declared earlier will return the first point of each sequence only, not all the points.
    # We could then iterate on the points using `next`, but since we only want the instruction's addresses and nothing
    # relevant to this instant in the trace (memory, cpu, etc), we will iterate on the basic_block's instructions instead.
    for head_point in reven_trace:
        bb = head_point.basic_block

        # Note we use the default BasicBlock's iterator
        for ins in bb:
            ins_addr = ins.address
            if ins_addr in bb_freq:
                bb_freq[ins_addr] += 1
    return {addr - relocation_offset: bb_freq[addr] for addr in bb_freq}


def collect_basic_blocks():
    bbs = set()
    for fun_head in idautils.Functions():
        fun_flowchart = idaapi.FlowChart(idaapi.get_func(fun_head))
        for bb in fun_flowchart:
            # In some cases, IDA's analysis gives incorrect results, e.g.
            #   1. startEA and endEA are identical
            if (bb.startEA != bb.endEA):
                bbs.add(bb)

    return bbs


def calculate_dynamic_basic_blocks(ida_bbs, reven_trace, relocation_offset):
    dynamic_basic_blocks = { (bb.startEA + relocation_offset, bb.endEA + relocation_offset) for bb in ida_bbs }

    for head_point in reven_trace:
        head_ins_addr = head_point.instruction.address
        try:
            (fst, snd) = next((fst, snd) for (fst, snd) in dynamic_basic_blocks if (head_ins_addr > fst) and (snd > head_ins_addr))

            dynamic_basic_blocks.remove((fst, snd))

            dynamic_basic_blocks.add((fst, head_ins_addr))
            
            dynamic_basic_blocks.add((head_ins_addr, snd))

        except StopIteration:
            pass
        except:
            pass

    dynamic_basic_blocks = { (fst - relocation_offset, snd - relocation_offset) for (fst, snd) in dynamic_basic_blocks }

    return dynamic_basic_blocks


def color_gradient(bb_freq):
    min_freq = bb_freq.itervalues().next()
    max_freq = min_freq
    for bb_addr in bb_freq:
        if min_freq > bb_freq[bb_addr]:
            min_freq = bb_freq[bb_addr]
        if max_freq < bb_freq[bb_addr]:
            max_freq = bb_freq[bb_addr]

        # calculate the green component of the basic block's color
        bb_color = {}
        max_color = 0xff
        if max_freq == min_freq:
            for bb_addr in bb_freq:
                bb_color[bb_addr] = max_color
        else:
            color_step = float(max_color) / (max_freq - min_freq)
            for bb_addr in bb_freq:
                bb_color[bb_addr] = max_color - int(round((bb_freq[bb_addr] - min_freq) * color_step))

    return bb_color


def colorize_dynamic_basic_block(bb, color):
    (addr, limit_addr) = bb
    while addr < limit_addr:
        idaapi.set_item_color(addr, color)
        addr = idc.NextHead(addr)


def colorize_dynamic_basic_blocks(dynamic_bbs, bb_freq):
    green_colors = color_gradient(bb_freq)

    bb_dict = { fst: (fst, snd) for (fst, snd) in dynamic_bbs }

    print '{:s}\t{:s}'.format('BAddr', 'Freq')

    for bb_fst_addr in bb_dict:
        if bb_freq[bb_fst_addr] > 0:
            # BBGGRR color format: BB = 0xb0, GG = green_colors[bb_head_addr], RR = 0xff
            color = 0xb0 * 0x10000 + green_colors[bb_fst_addr] * 0x100 + 0xff
            colorize_dynamic_basic_block(bb_dict[bb_fst_addr], color)

            print '0x{:x}\t{:d}'.format(bb_fst_addr, bb_freq[bb_fst_addr])


if __name__ == '__main__':
    host_port_str = idc.AskStr('localhost:13370', "REVEN's project address")
    if host_port_str is not None:
        try:
            host, port_str = host_port_str.split(':')
            port = int(port_str)
            print("REVEN's project: {}:{}").format(host, port)
            main(host, port)
        except ValueError:
            print("please give a correct REVEN\'s project address, e.g. localhost:13370")
        except RuntimeError, e:
            print('{}').format(e)
        except:
            print('Unknown error')