r2-call-graph.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
------------------------------------------------------------------------------
Script Name:   r2-call-graph.py
Description:   Generates a TARGETED call graph rooted at a specific function,
               with a configurable maximum depth. Complements
               `r2-dissector.py --call-graph`, which produces the GLOBAL call
               graph (every function in the binary).

               Use this script when the global graph is too dense to read or
               when you want a focused view of "what does main() reach in
               three levels?".

               v1.1 changes:
                 - Cleaned up name resolution when an address (hex/int) is
                   passed as the start function.
                 - Added complexity-based color legend.
                 - Allows redirecting output to a custom directory.

Author:        Generated by Google Gemini, Revised by Ricardo J. Rodríguez
Created:       2025-12-14
Version:       1.1

Requirements:  pip install r2pipe graphviz
               System required: Graphviz (`dot`) and radare2 in PATH.
------------------------------------------------------------------------------
"""

import r2pipe
import sys
import os
import argparse
from graphviz import Digraph


def log(msg):
    sys.stderr.write(f"[+] {msg}\n")


def error(msg):
    sys.stderr.write(f"[-] {msg}\n")


def safe_id(name):
    """Sanitises a string for use as a Graphviz node ID."""
    if not name:
        return "unknown"
    out = str(name)
    for ch in '.:@ -':
        out = out.replace(ch, '_')
    return out


def resolve_start_name(token, func_meta):
    """If `token` looks like an address, try to map it to its symbol name."""
    if isinstance(token, int):
        meta = func_meta.get(token)
        return meta['name'] if meta else hex(token)

    if isinstance(token, str) and token.startswith("0x"):
        try:
            addr = int(token, 16)
            meta = func_meta.get(addr)
            if meta:
                return meta['name']
        except ValueError:
            pass

    return token


def color_for_cc(cc):
    if cc > 50:
        return "red"
    if cc > 20:
        return "orange"
    if cc > 10:
        return "yellow"
    return "lightblue"


def generate_call_graph(binary_path, output_pdf, start_func=None, max_depth=5, min_cc=0):
    log(f"Opening {binary_path}...")
    r2 = r2pipe.open(binary_path, flags=['-2'])

    log("Analyzing binary (aaa)...")
    r2.cmd("aaa")

    log("Fetching function list...")
    all_funcs = r2.cmdj("aflj") or []

    # Lookup by name AND by address.
    func_meta = {}
    for f in all_funcs:
        name = f.get('name')
        addr = f.get('offset')
        meta = {'cc': f.get('cc', 0), 'name': name, 'addr': addr}
        if name:
            func_meta[name] = meta
        if addr is not None:
            func_meta[addr] = meta

    dot = Digraph(comment='Call Graph', strict=True)
    dot.attr(rankdir='LR')
    dot.attr('node', shape='box', style='filled', color='lightblue', fontname='Helvetica')

    nodes_to_draw = set()
    edges_to_draw = set()
    visited = set()

    if start_func:
        start_func = resolve_start_name(start_func, func_meta)
        log(f"Mode: targeted tracing from '{start_func}' (max depth: {max_depth})")

        queue = [(start_func, 0)]
        while queue:
            curr_name, depth = queue.pop(0)
            if curr_name in visited or depth > max_depth:
                continue

            visited.add(curr_name)
            nodes_to_draw.add(curr_name)

            # Only known functions can be queried with axffj; raw addresses
            # and imports just become leaves (no further recursion).
            if curr_name not in func_meta:
                continue

            refs = r2.cmdj(f"axffj @ {curr_name}") or []
            for ref in refs:
                if ref.get('type') not in ('CALL', 'JMP', 'CODE'):
                    continue
                target_name = ref.get('name') or f"0x{ref.get('to', 0):x}"
                edges_to_draw.add((curr_name, target_name))
                if target_name not in visited:
                    queue.append((target_name, depth + 1))
                    nodes_to_draw.add(target_name)

    else:
        log(f"No start function specified. Graphing ALL (filtered by min_cc={min_cc})...")
        for f in all_funcs:
            if f.get('cc', 0) >= min_cc:
                fname = f.get('name')
                if fname:
                    nodes_to_draw.add(fname)
                    visited.add(fname)

        for fname in list(nodes_to_draw):
            refs = r2.cmdj(f"axffj @ {fname}") or []
            for ref in refs:
                if ref.get('type') not in ('CALL', 'JMP'):
                    continue
                tname = ref.get('name') or f"0x{ref.get('to', 0):x}"
                edges_to_draw.add((fname, tname))
                nodes_to_draw.add(tname)

    log(f"Generating graph with {len(nodes_to_draw)} nodes / {len(edges_to_draw)} edges...")

    for fname in nodes_to_draw:
        cc = func_meta.get(fname, {}).get('cc', 0)
        color = color_for_cc(cc)
        if start_func and fname == start_func:
            color = "lightgreen"
        dot.node(safe_id(fname), label=f"{fname}\nCC: {cc}", fillcolor=color)

    for src, dst in edges_to_draw:
        dot.edge(safe_id(src), safe_id(dst))

    # Legend (small subgraph at the bottom-right).
    with dot.subgraph(name='cluster_legend') as lg:
        lg.attr(label='CC legend', fontsize='10', style='dashed')
        for label, cc in [("CC ≤ 10", 0), ("CC > 10", 11), ("CC > 20", 21), ("CC > 50", 60)]:
            lg.node(f"_legend_{cc}", label=label, fillcolor=color_for_cc(cc), style='filled')

    output_filename = os.path.splitext(output_pdf)[0]
    try:
        dot.render(output_filename, view=False, format='pdf', cleanup=True)
        log(f"Saved: {output_filename}.pdf")
    except Exception as e:
        error(f"Graphviz error: {e}")
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        description="Targeted call graph rooted at a specific function.",
    )
    parser.add_argument("binary", help="Path to the binary file")
    parser.add_argument("output", help="Output PDF filename")
    parser.add_argument("-s", "--start", type=str, default=None,
                        help="Root function (name or hex address, e.g. main, 0x401000)")
    parser.add_argument("-d", "--depth", type=int, default=3,
                        help="Maximum recursion depth from the start function (default: 3)")
    parser.add_argument("-c", "--min-cc", type=int, default=0,
                        help="Minimum cyclomatic complexity (only used when --start is unset)")
    args = parser.parse_args()

    if not os.path.exists(args.binary):
        error(f"File not found: {args.binary}")
        sys.exit(1)

    generate_call_graph(args.binary, args.output, args.start, args.depth, args.min_cc)


if __name__ == "__main__":
    main()