#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
#
# callwalk.py - walk the Open-IOV callgraph index and print callee trees /
# caller trees around a function. This is the quick way to go from an export
# wrapper to the internal implementation chain without opening Ghidra.
#
# Usage:
#   callwalk.py <callgraph.tsv> callees <hex_file_vaddr> [depth] [max_per_node]
#   callwalk.py <callgraph.tsv> callers <hex_file_vaddr> [depth] [max_per_node]
#   callwalk.py <callgraph.tsv> chain   <hex_from> <hex_to>   # shortest path
#
# Addresses are FILE virtual addresses (libcuda+0x...); the callgraph columns
# `caller_file`/`callee_file` are used, so Ghidra image-base offsets do not
# matter. `names.tsv` (file_vaddr<TAB>name) is picked up next to the callgraph
# when present, and used to label nodes.
#
# SPDX-License-Identifier: BSD-3-Clause

import sys
from collections import defaultdict, deque


def load_graph(path):
    callees = defaultdict(list)
    callers = defaultdict(list)
    with open(path) as f:
        header = f.readline()
        for line in f:
            parts = line.rstrip("\n").split("\t")
            if len(parts) < 4:
                continue
            cf, kf = int(parts[2], 16), int(parts[3], 16)
            callees[cf].append(kf)
            callers[kf].append(cf)
    return callees, callers


def load_names(path):
    names = {}
    try:
        with open(path) as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                parts = line.split("\t")
                if len(parts) >= 2:
                    try:
                        names[int(parts[0], 16)] = parts[1]
                    except ValueError:
                        pass
    except FileNotFoundError:
        pass
    return names


def label(a, names):
    return f"0x{a:08x}({names[a]})" if a in names else f"0x{a:08x}"


def walk(adj, start, depth, limit):
    seen = {start}
    frontier = [start]
    for d in range(depth):
        nxt = []
        for node in frontier:
            kids = sorted(set(adj.get(node, [])))
            if not kids:
                continue
            print("  " * (d + 1) + " | ".join(label(k, NAMES) for k in kids[:limit])
                  + (" ..." if len(kids) > limit else ""))
            for k in kids:
                if k not in seen:
                    seen.add(k)
                    nxt.append(k)
        frontier = nxt
        if not frontier:
            break


def shortest_path(callees, start, goal):
    if start == goal:
        return [start]
    prev = {start: None}
    q = deque([start])
    while q:
        node = q.popleft()
        for k in callees.get(node, []):
            if k in prev:
                continue
            prev[k] = node
            if k == goal:
                path = [goal]
                while prev[path[-1]] is not None:
                    path.append(prev[path[-1]])
                return list(reversed(path))
            q.append(k)
    return None


def main():
    if len(sys.argv) < 4:
        print(__doc__)
        sys.exit(2)
    graph = sys.argv[1]
    mode = sys.argv[2]
    addr = int(sys.argv[3], 16)
    depth = int(sys.argv[4]) if len(sys.argv) > 4 else 4
    limit = int(sys.argv[5]) if len(sys.argv) > 5 else 8

    global NAMES
    callees, callers = load_graph(graph)
    NAMES = load_names(graph.rsplit("/", 1)[0] + "/names.tsv")

    if mode == "chain":
        goal = int(sys.argv[4], 16)
        path = shortest_path(callees, addr, goal)
        if not path:
            print("no direct path found")
        else:
            for a in path:
                print(label(a, NAMES))
        return

    print(f"{mode} of {label(addr, NAMES)}:")
    walk(callees if mode == "callees" else callers, addr, depth, limit)


NAMES = {}

if __name__ == "__main__":
    main()