ida_xref_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_xref_simple.py
  • x86: C:\Program Files (x86)\Reven-PythonAPI-x86\python-examples\ida\ida_xref_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 to add code cross references for dynamic
    control flow instructions (whose targets cannot be computed by IDA's static
    analysis) of a binary in IDA. The cross references are local, namely callees
    being not in the binary will not be counted.

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:
    The script takes targets from the execution trace logged by REVEN (i.e. dynamic
    analysis). It has the following limits:
        - 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

        - the added cross-refs are not complete (as an inherent problem of dynamic
          analysis)

        - self-modifying/overlapping instructions are checked in a very limited way
          (does not assume any coherent result in case of self-modifying/packed code).
"""

import idc
import idautils
import idaapi

import reven


ida_dynamic_jump_types = [
    idaapi.NN_jmp,
    idaapi.NN_jmpfi,
    idaapi.NN_jmpni,
    idaapi.NN_jmpshort
]

ida_dynamic_call_types = [
    idaapi.NN_call,
    idaapi.NN_callfi,
    idaapi.NN_callni
]

ida_dynamic_ret_types = [
    idaapi.NN_retn,
    idaapi.NN_retf
]

reven_jcc_mnemonics = ['jmp', 'call', 'retn', 'retf']


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)

        inss = collect_indirect_jccs()
        trace = get_binary_trace(project)
        offset = runtime_base_address - static_base_address

        add_xrefs(inss, trace, offset)

        print 'adding code xref 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 reven_points


def add_xrefs(ida_jc_inss, reven_points, relocation_offset):
    jcc_dict = {ins.ea: idc.GetManyBytes(ins.ea, ins.size) for ins in ida_jc_inss}

    bb_control_flow = []
    examined_control_flow = set()
    prev_point = None

    # Let's iterate on the binary's trace and store relevant jumps in bb_control_flow.
    # While doing so, we will have to filter out edge cases such as the binary jumping out to another.
    # Also, the generator we declared earlier will return the first point of each sequence only, not all the points, so
    # getting the information we want requires a bit of fiddling.
    for curr_point in reven_points:
        if prev_point is not None:
            # Look for the last instruction of the previous basic block and
            # its next instruction
            target_point = prev_point
            for last_ins_of_prev_bb in prev_point.basic_block:
                target_point = target_point.next()

            # The target point is always the next instruction, but it is not always
            # in the local trace (e.g. it may be the target of a call to an
            # external function)
            curr_ins = curr_point.instruction
            target_ins = target_point.instruction

            # We have to check:
            #   1. the target instruction is in the local trace
            #   2. the caller is a dynamic control flow instruction
            #   3. the caller have been discovered by the static analysis
            #   4. the control flow (caller, callee) is not counted
            if (curr_ins.address == target_ins.address) and \
               (last_ins_of_prev_bb.mnemonic in reven_jcc_mnemonics) and \
               (last_ins_of_prev_bb.address - relocation_offset in jcc_dict) and \
               ((last_ins_of_prev_bb.address, curr_ins.address) not in examined_control_flow):
                examined_control_flow.add((last_ins_of_prev_bb.address, curr_ins.address))
                bb_control_flow.append((last_ins_of_prev_bb, curr_ins))

        prev_point = curr_point

    # Now bb_control_flow should contain a list of (caller, callee) extracted from
    # the REVEN's project trace

    updated_jmp_xrefs = set()
    updated_call_xrefs = set()
    updated_ret_xrefs = set()

    for (prev_ins, curr_ins) in bb_control_flow:
        mnemonic = prev_ins.mnemonic
        caller_address = prev_ins.address - relocation_offset

        if caller_address in jcc_dict:
            if jcc_dict[caller_address] != prev_ins.raw_bytes:
                print 'warning: instruction at 0x{:x} is modified in running or loading time'.format(caller_address)

            target_address = curr_ins.address - relocation_offset
            if mnemonic == 'jmp':
                if (caller_address, target_address) not in updated_jmp_xrefs:
                    idc.AddCodeXref(caller_address, target_address, idc.fl_JF)
                    updated_jmp_xrefs.add((caller_address, target_address))
            elif mnemonic == 'call':
                if (caller_address, target_address) not in updated_call_xrefs:
                    idc.AddCodeXref(caller_address, target_address, idc.fl_CF)
                    updated_call_xrefs.add((caller_address, target_address))
            else: # mnemonic == 'ret'
                if (caller_address, target_address) not in updated_ret_xrefs:
                    idc.AddCodeXref(caller_address, target_address, idc.fl_JF)
                    updated_ret_xrefs.add((caller_address, target_address))

    print 'updated xrefs from runtime information:'
    print '  jmp:  {:d}'.format(len(updated_jmp_xrefs))
    print '  call: {:d}'.format(len(updated_call_xrefs))
    print '  ret:  {:d}'.format(len(updated_ret_xrefs))

    for (caller_address, target_address) in updated_jmp_xrefs:
        print '0x{:x} => 0x{:x} (jmp)'.format(caller_address, target_address)

    for (caller_address, target_address) in updated_call_xrefs:
        print '0x{:x} => 0x{:x} (call)'.format(caller_address, target_address)

    for (caller_address, target_address) in updated_ret_xrefs:
        print '0x{:x} => 0x{:x} (ret)'.format(caller_address, target_address)


def collect_indirect_jccs():
    indirect_calls = []
    indirect_jumps = []
    indirect_rets = []

    for seg in idautils.Segments():
        for head in idautils.Heads(idc.SegStart(seg), idc.SegEnd(seg)):
            if idc.isCode(idc.GetFlags(head)):
                ins = idautils.DecodeInstruction(head)
                if ins is not None:
                    if ins.get_canon_feature() & idaapi.CF_JUMP:
                        if ins.itype in ida_dynamic_jump_types:
                            if ins not in indirect_jumps:
                                indirect_jumps.append(ins)
                        elif ins.itype in ida_dynamic_call_types:
                            if ins not in indirect_calls:
                                indirect_calls.append(ins)
                    elif ins.itype in ida_dynamic_ret_types:
                        if ins not in indirect_rets:
                                indirect_rets.append(ins)

    print 'statically detected jccs:'
    print '  jmp:  {:d}'.format(len(indirect_jumps))
    print '  call: {:d}'.format(len(indirect_calls))
    print '  ret:  {:d}'.format(len(indirect_rets))

    return indirect_calls + indirect_jumps + indirect_rets


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')