from binaryninja import *
from binaryninjaui import UIAction, UIActionHandler, Menu
from PySide6.QtGui import QKeySequence


def main(ctx):
    from binaryninja.lowlevelil import LowLevelILInstruction, LowLevelILOperation
    from binaryninja import log_info

    def get_curr_llil_ssa(ctx):
        """Retrieves the currently selected (UI) IL instruction"""
        func = ctx.function
        view = ctx.view
        if not func:
            log_info("Not inside a function")
            return None
        if view.getILViewType().view_type != FunctionGraphType.LowLevelILSSAFormFunctionGraph:
            log_info("Function is not in llil ssa mode")
            return None

        # Get the currently selected instruction (UI)
        ssa_idx = view.getCurrentILInstructionIndex()
        if ssa_idx is None:
            log_info("Failed to retrieve currently selected inst index")
            return None
        llil_ssa_func = func.low_level_il.ssa_form
        inst = llil_ssa_func[ssa_idx]
        return inst

    def visit(graph, parent_node, expr, ops):
        """Recursively visit each expression and build graph of the expression tree"""
        if expr is None:
            return
        op = expr.operation

        # Build graph node
        child_node = FlowGraphNode(graph)
        child_node.lines = [f"{op.name}"]
        graph.append(child_node)
        # If this is the first node in the tree it has no parent
        if parent_node is not None:
            parent_node.add_outgoing_edge(BranchType.UnconditionalBranch, child_node)

        # Recurse through children
        for name, kind in ops[op]:
            print(f"{kind}: {name}")
            if kind == "expr":
                visit(graph, child_node, getattr(expr, name), ops)
            elif kind == "expr_list":
                for sub in getattr(expr, name):
                    visit(graph, child_node, sub, ops)

    il_ops = LowLevelILInstruction.ILOperations
    inst = get_curr_llil_ssa(ctx)
    if inst is None:
        return

    graph = FlowGraph()
    visit(graph, None, inst, il_ops)
    # This spawns a new tab with our little graph in it
    show_graph_report("Expression tree", graph)


PluginCommand.register("SSA Graphing", "Graphs an LLIL SSA AST", main)
# Register a keybind for the analysis
UIAction.registerAction("Graph LLIL SSA AST", QKeySequence("Ctrl + F3"))
UIActionHandler.globalActions().bindAction("Graph LLIL SSA AST", UIAction(main))

Menu.mainMenu("Plugins").addAction("Graph LLIL SSA AST", "Plugins")