#!/usr/bin/env python

import r2pipe
from anytree import Node, RenderTree

r2 = r2pipe.open()

here = int(r2.cmd("s"), 16)

print "Here: ", hex(here)

history = {}

functions = r2.cmdj('aflj')


# There is probably a better way to do this
def getResidingFunction(addr):
    for func in functions:
        if addr >= func['offset'] and addr < func['offset'] + func['realsz']:
            return func
    return None


def printTree(root_node):
    for pre, _, node in RenderTree(root_node):
        str_loc = "{0:#0{1}x}".format(node.name, 10)
        fcn_name = r2.cmd("afl~" + str_loc + "[3]")
        if fcn_name == "->":
            fcn_name = r2.cmd("afl~" + str_loc + "[5]")
        print("%s%s - %s" % (pre, str_loc, fcn_name))


def populateTree(current_node):
    xrefs = r2.cmdj("axtj @ " + hex(current_node.name))

    if len(xrefs) == 0:
        # Leaf detected
        return
    else:
        for xref in xrefs:
            fcn = getResidingFunction(int(xref['from']))
            if fcn != None and int(fcn['offset']) != current_node.name:
                # Found an xref from inside a function
                new_node = Node(int(fcn['offset']))
            else:
                # Found an xref from outside function bounds
                new_node = Node(xref['from'])

            # Check if we've processed this node before
            try:
                if history[new_node.name] != None:
                    # Merging paths detected
                    history[current_node.name] = 1
                    continue
            except:
                pass
            history[current_node.name] = 1

            current_node.children += (new_node,)
            populateTree(new_node)

        return

here_node = Node(here)
populateTree(here_node)
printTree(here_node)