cfg_from_trace.py

• Linux path to script:
  • /usr/share/reven/examples/cfg_from_trace.py
• Windows path to script:
  • amd64: C:\Program Files (x86)\Reven-PythonAPI-amd64\python-examples\ida\cfg_from_trace.py
  • x86: C:\Program Files (x86)\Reven-PythonAPI-x86\python-examples\ida\cfg_from_trace.py

This example demonstrates how to use the python API to build the CFG from the trace of a given process.

See the script's documentation below for more information.

Note that this script only iterates over points that belong to a given process, and not over the whole trace. For more details, look at the implementation of the get_process_point_ranges and get_process_sequence_points methods.

"""
Purpose:
  This script builds the control flow graph from the trace of a given process.

Usage:
  Give the following arguments:
   - Host of the REVEN's project. (default: "localhost")
   - Port of the REVEN's project. (default = 13370)
   - Process CR3 or PID.
   - Output filename.

  For example:
   python cfg_from_trace.py --host localhost --port 13370 cr3 0x1e7e3000 output_dot_filename

Remark: The self-modifying/overlapping code is not supported
"""

import reven
import pygraphviz
import sys
import argparse

def check_process(project, cr3):
    """
    Check if a process is running in the whole trace.
    @param cr3: The CR3 register's value associated to the process to check.
    """
    for process in project.processes():
        if process.cr3 == cr3:
            return
    raise ValueError('Unknown process with cr3 = {}'.format(hex(cr3)))

def get_cr3_from_pid(project, pid):
    """
    Get the CR3 register's value associated to a running process.
    @param pid: The PID of the process.
    """
    for process in project.processes():
        if process.pid == pid:
            return process.cr3
    raise ValueError('Unknown process with pid = {}'.format(pid))

def get_last_point(trace):
    """
    Get the last point of the given trace.
    """
    last_sequence = trace.sequence_count - 1
    last_instruction = trace.sequence_length(last_sequence) - 1
    return trace.point(last_sequence, last_instruction)

def get_process_point_ranges(project, cr3):
    """
    Get the ranges of main trace (e.g. 'Execution run') points where the process is running.
    @param cr3: The CR3 register's value.
    @return: A generator of tuples composed of the start and end point([start, end]).
    """
    main_trace = project.traces()[0] # Execution run
    first_point = main_trace.point(0, 0)
    start = first_point if first_point.cpu().read_register('cr3') == cr3 else None
    for switch in project.process_switches():
        if main_trace.name != switch.point.trace.name:
            continue
        if start is None:
            if cr3 == switch.cr3:
                start = switch.point.next_sequence()
        else:
            yield (start, switch.point)
            start = None
    if start:
        # special case where the trace finishes while the process is running.
        yield (start, get_last_point(main_trace))

def get_process_number_of_sequences(project, cr3):
    process_point_ranges = get_process_point_ranges(project, cr3)
    return sum([ last_point.sequence_index - first_point.sequence_index + 1 \
                 for first_point, last_point in process_point_ranges ])

def get_process_sequence_points(project, cr3):
    """
    Get the first points of main trace (e.g. 'Execution run') sequences where the process is running in the main trace.
    @param cr3: The CR3 register's value.
    @return: A generator of points belonging to the main trace.
    """
    for first, last in get_process_point_ranges(project, cr3):
        point = first
        while point is not None and point <= last:
            yield point
            point = point.next_sequence()

class EmptySequencePoints(RuntimeError):
    pass

def construct_control_flow(head_points, number_of_head_points):
    if number_of_head_points == 0:
        raise EmptySequencePoints('The process is not running in the main trace.')

    percentage_step_length = number_of_head_points / 100
    percentage_step = 0
    examined_head_point_number = 0

    print('Number of basic blocks in the trace: {}').format(number_of_head_points)
    print('Building control flow graph, please wait')

    # CFG building algorithm (c.f. EAC2e, c.5.3.4)
    distinguished_instructions = {}
    leaders = set()

    entry_point = None
    previous_follower = None
    followers = set()
    control_flow = set()

    # first pass: detect leaders
    for point in head_points:
        has_head_point = True
        leader = point.instruction.address

        if previous_follower is not None:
            control_flow.add((previous_follower, leader))
        else:
            entry_point = leader

        leaders.add(leader)

        reven_basic_block = point.basic_block

        for ins in reven_basic_block:
            if ins.address not in distinguished_instructions:
                distinguished_instructions[ins.address] = ins

        previous_follower = ins.address
        followers.add(previous_follower)

        examined_head_point_number += 1
        percentage_step += 1
        if percentage_step >= percentage_step_length:
            percentage_step = 0
            print('...{:0.3f}%').format(examined_head_point_number * 100.0 / number_of_head_points),
            sys.stdout.flush()

    print('...completed')

    # second pass: build basic blocks
    basic_blocks = {}
    ins_iter = iter(sorted(distinguished_instructions.iterkeys()))
    try:
        ins = next(ins_iter)
        while True:
            if ins not in leaders:
                ins = next(ins_iter)
            else:
                leader = ins
                basic_blocks[leader] = []

                while True:
                    prev_ins = ins
                    basic_blocks[leader].append(prev_ins)
                    ins = next(ins_iter)

                    if (prev_ins in followers) or (ins in leaders):
                        if prev_ins not in followers: # then ins must be a leader
                            followers.add(prev_ins)
                            control_flow.add((prev_ins, ins))
                        break
    except StopIteration:
        pass

    # different from the original algorithm which works always given static form
    # of the program, we need a third pass to build control flow since we cannot
    # be sure that basic blocks are consecutive
    basic_block_control_flow = set()
    for leader in basic_blocks:
        follower = basic_blocks[leader][-1]
        for head, tail in control_flow:
            if follower == head:
                basic_block_control_flow.add((leader, tail))

    return (distinguished_instructions, entry_point, leaders, basic_blocks, basic_block_control_flow)


def build_control_flow_graph(instructions, entry_point, leaders, basic_blocks, control_flow):
    def instruction_to_string(reven_ins):
        ins_mnemonic = reven_ins.mnemonic
        ins_operands = ', '.join(reven_ins.operands())
        return ('0x{:<10x}{:<8}{}'.format(reven_ins.address, ins_mnemonic, ins_operands)).lower()

    def build_node_label(leader):
        basic_block = [instructions[address] for address in basic_blocks[leader]]
        return ('\l'.join([instruction_to_string(ins) for ins in basic_block])) + '\l'

    cfg = pygraphviz.AGraph(strict=False, directed=True, name='Control Flow Graph')
    cfg.node_attr['shape'] ='box'
    cfg.node_attr['style'] = 'rounded'
    cfg.node_attr['fontname'] = 'Liberation Mono'

    for leader in leaders:
        if leader == entry_point:
            cfg.add_node(hex(leader), 
                         label='entry point\n' + build_node_label(leader),
                         style='filled, rounded',
                         fillcolor='cornflowerblue')
        else:
            cfg.add_node(hex(leader),
                         label=build_node_label(leader))

    for head, tail in control_flow:
        cfg.add_edge(hex(head), hex(tail))

    return cfg


def main(reven_project, cr3, pid, output_graph_filename):
    if cr3 is not None:
        check_process(reven_project, cr3)
    else:
        cr3 = get_cr3_from_pid(reven_project, pid)

    head_points = get_process_sequence_points(reven_project, cr3)
    number_of_head_points = get_process_number_of_sequences(reven_project, cr3)

    instructions, entry_point, leaders, basic_blocks, control_flow = construct_control_flow(head_points, number_of_head_points)

    print('Entry point: 0x{:x}').format(entry_point)
    print('Instructions: {:d}').format(len(instructions))
    print('Leaders: {:d}').format(len(leaders))
    print('Control flow: {:d}').format(len(control_flow))

    control_flow_graph = build_control_flow_graph(instructions, entry_point, leaders, basic_blocks, control_flow)
    control_flow_graph.draw(path=output_graph_filename, format='dot', prog='dot')
    print('Control flow graph built, please check the output dot file: {}').format(output_graph_filename)

def auto_int(x):
    return int(x, 0)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, help='The reven server host. (default = "localhost")', default="localhost")
    parser.add_argument("--port", type=int, help="The reven server port. (default = 13370)", default=13370)
    subparsers = parser.add_subparsers(dest="subparser", help="Process identification")
    cr3_parser = subparsers.add_parser('cr3', help="Use the CR3 register to identify a process (recommended).")
    cr3_parser.add_argument("value", type=auto_int, help="The CR3 register value")
    pid_parser = subparsers.add_parser('pid', help="Use the PID to identify a process.")
    pid_parser.add_argument("value", type=int, help="The PID value")
    parser.add_argument("out", type=str, help="The output filename (without extension).")
    args = parser.parse_args()

    cr3 = args.value if args.subparser == 'cr3' else None
    pid = args.value if args.subparser == 'pid' else None

    try:
        project = reven.Project(args.host, args.port)
        main(project, cr3, pid, args.out + '.dot')
    except (RuntimeError, ValueError, EmptySequencePoints) as e:
        print('Error: {}').format(e)
    except:
        print('Unknown error: {}').format(sys.exc_info())