Skip to content

Instantly share code, notes, and snippets.

@ophirharpaz
Created February 22, 2020 15:19
Show Gist options
  • Save ophirharpaz/c7f6314006ad33a345379d186f51b2d7 to your computer and use it in GitHub Desktop.
Save ophirharpaz/c7f6314006ad33a345379d186f51b2d7 to your computer and use it in GitHub Desktop.
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.
"""
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())
@ophirharpaz
Copy link
Author

get_call_flows_from_exports.py

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment