The script generates and prints a graph of all function-call flows that start in exported functions and end in the function being pointed at in IDA. This functionality is useful when you need to trigger a function in a DLL and wish to know which exported function leads to it.

get_call_flows_from_exports.py

"""
The script generates and prints a graph of all function-call flows that start in exported functions and end
in the function being pointed at in IDA.
This functionality is useful when you need to trigger a function in a DLL and wish to know which exported function
leads to it.
"""

import idaapi
import idautils
import idc
import networkx as nx
import sark

EXPORTS = [e[2] for e in idautils.Entries()]  # entry[2] stores its address
ROOT_NODE_COLOR = 0xA67111
TARGET_NODE_COLOR = 0x5f1299


class NameNodeHandler(sark.ui.AddressNodeHandler):
    """
    Inheritance from AddressNodeHandler gives the already-implemented on_click functionality.
    """

    def on_get_text(self, value, attrs):
        function_name = idc.GetFunctionName(value)
        demangled = idc.Demangle(function_name, idc.GetLongPrm(idc.INF_SHORT_DN))
        if demangled:
            return demangled.split('(')[0]
        return function_name


def is_xref_a_function_call(xref):
    return xref.type in [idaapi.fl_CF, idaapi.fl_CN]


def get_calling_functions(function_address):
    # Returns the addresses of all functions that call function_address
    xrefs = (xref.frm for xref in idautils.XrefsTo(function_address) if is_xref_a_function_call(xref))
    # Some xrefs cannot be "attributed" to a function - skip those
    referencing_functions = {idaapi.get_func(xref).startEA for xref in xrefs if idaapi.get_func(xref)}
    return referencing_functions


def _rec_get_call_sub_graph(function_address, call_sub_graph, visited_functions, is_reachable_from_exports):
    if function_address in EXPORTS:
        is_reachable_from_exports[function_address] = True
    
    # If the function was visited before during recursion, return its verdict
    if function_address in visited_functions:
        return is_reachable_from_exports[function_address]
    
    calling_functions = get_calling_functions(function_address)  # All referencing functions
    
    # Mark the current function as "visited" and set it to False by default (until found otherwise)
    visited_functions.add(function_address)
    is_reachable_from_exports[function_address] = False
    
    # Relevant caller functions are ones which:
    # (a) haven't been visited yet; OR
    # (b) have been already found to be reachable by exported functions.
    relevant_callers = (calling_functions
                        - visited_functions
                        - {k for k, v in is_reachable_from_exports.items() if not v})
    
    for caller in relevant_callers:
        edge = (caller, function_address)
        # Add the edge if it's not already in the graph and if the caller is reachable by an exported function
        if (edge not in call_sub_graph.edges
                and _rec_get_call_sub_graph(caller, call_sub_graph, visited_functions, is_reachable_from_exports)):
            call_sub_graph.add_edge(*edge)
            is_reachable_from_exports[function_address] = True
    
    return is_reachable_from_exports[function_address]


def get_call_sub_graph(function_address):
    call_graph = nx.DiGraph()
    visited_functions = set()
    is_reachable_from_exports = {}
    _rec_get_call_sub_graph(function_address, call_graph, visited_functions, is_reachable_from_exports)
    return call_graph


def print_call_sub_graph(function):
    # (1) Generate the sub-graph of all function-calls that lead to function
    function_address = function.startEA
    call_graph = get_call_sub_graph(function_address)

    # (2) Draw the graph in IDA
    if call_graph:
        # Reverse the graph such as edges point to our function
        call_graph.reverse()

        # Colorize root & target nodes
        call_graph.node[function_address][sark.ui.NXGraph.BG_COLOR] = TARGET_NODE_COLOR
        root_nodes = set(EXPORTS).intersection(call_graph.nodes())
        for node_address in root_nodes:
            call_graph.node[node_address][sark.ui.NXGraph.BG_COLOR] = ROOT_NODE_COLOR

        viewer = sark.ui.NXGraph(call_graph, handler=NameNodeHandler())
        viewer.Show()
    else:
        print('Function {} cannot be reached from an export'.format(function.name))


if __name__ == '__main__':
    print_call_sub_graph(sark.Function())