IDA Graph view with outlined function included

outline_graph.py

"""
summary: drawing custom graphs
description:
  Showing custom graphs, using `ida_graph.GraphViewer`. In addition,
  show how to write actions that can be performed on those.
keywords: graph, actions
"""

from __future__ import print_function
# -----------------------------------------------------------------------
# This is an example illustrating how to use the user graphing functionality
# in Python
# (c) Hex-Rays
#

import ida_kernwin
import ida_graph
import ida_ua
import ida_idp
import ida_funcs
import ida_xref

import idautils

class _base_graph_action_handler_t(ida_kernwin.action_handler_t):
    def __init__(self, graph):
        ida_kernwin.action_handler_t.__init__(self)
        self.graph = graph

    def update(self, ctx):
        return ida_kernwin.AST_ENABLE_ALWAYS


class GraphCloser(_base_graph_action_handler_t):
    def activate(self, ctx):
        self.graph.Close()


class SelectionPrinter(_base_graph_action_handler_t):
    def activate(self, ctx):
        try:
            sel = ctx.graph_selection
        except:
            # IDA < 7.4 doesn't provide graph selection as part of
            # the action_activation_ctx_t; it needs to be queried.
            sel = ida_graph.screen_graph_selection_t()
            gv = ida_graph.get_graph_viewer(self.graph.GetWidget())
            ida_graph.viewer_get_selection(gv, sel)
        if sel:
            for s in sel:
                if s.is_node:
                    print("Selected node %d" % s.node)
                else:
                    print("Selected edge %d -> %d" % (s.elp.e.src, s.elp.e.dst))
        return 1

import time
class MyGraph(ida_graph.GraphViewer):
    def __init__(self, func, chart):
        self.title = "call graph of %x %d" % (func.start_ea, time.time())
        ida_graph.GraphViewer.__init__(self, self.title)
        self.func = func
        self.chart = chart
        self.color = 0xff00ff
        self.nodes = {}
        self.Clear()

    def OnRefresh(self):
        print('OnRefresh')
        self.Clear()
        self.nodes = {}
        funcChain = []
        def add_funcchart(func, funcChain):
            if func.start_ea in funcChain:
                raise Exception('loop in outline func!')
            funcChain = funcChain + (func.start_ea,)
            
            q = idaapi.qflow_chart_t("The title", func, 0, 0, idaapi.FC_CALL_ENDS)
            qblks = [
                (
                    blki,
                    q[blki].start_ea, q[blki].end_ea, 
                    [ q.succ(blki, j) for j in range(q.nsucc(blki)) ],
                    [ q.pred(blki, j) for j in range(q.npred(blki)) ]
                )
                for blki in range(q.size())
            ]
            
            for blk in qblks:
                blki, blk_start, blk_end, blk_succ, blk_pred = blk
                if blk_start != blk_end:
                    continue
                qblks[blki] = None
                for i in range(len(qblks)):
                    cblk = qblks[i]
                    if not cblk:
                        continue
                    if blki in cblk[-2]:
                        cblk[-2].remove(blki)
                    if blki in blk[-1]:
                        cblk[-1].remove(blki)
            
            outlink_blks = {}
            for blk in qblks:
                if blk is None:
                    continue
                blki, blk_start, blk_end, blk_succ, blk_pred = blk
                lastins = idaapi.prev_head(blk_end, 0)
                xrefs = [x for x in XrefsFrom(lastins) if x.type in (fl_CN, fl_CF, fl_JN, fl_JF)]
                outlinef = None
                for x in xrefs:
                    xf = idaapi.get_func(x.to)
                    if xf and xf.start_ea != func.start_ea:
                        if xf.flags & idaapi.FUNC_OUTLINE:
                            outlinef = xf
                            break
                if outlinef:
                    print("Outline blk: %d %x" % (blki, outlinef.start_ea))
                    outlink_blks[blki] = outlinef
            
            for blk in reversed(qblks): 
                if blk is None:
                    continue
                blki, blk_start, blk_end, blk_succ, blk_pred = blk
                
                if len(blk_pred) == 1:
                    prev_blki, prevblk_start, prevblk_end, prevblk_succ, prevblk_pred = qblks[blk_pred[0]]
                    print(hex(blk_start), blki, blk_pred,prevblk_succ)
                    if prev_blki in outlink_blks:
                        continue
                    if blki in outlink_blks:
                        outlink_blks[prev_blki] = outlink_blks[blki]
                    if len(prevblk_succ) == 1 and prevblk_succ[0] == blki:
                        if prevblk_end == blk_start:
                            # need merge to prev
                            print('merge block %d' % blki)
                            qblks[prev_blki] = (
                                prev_blki,
                                prevblk_start, blk_end, 
                                blk_succ,
                                prevblk_pred
                            )
                            qblks[blki] = None
        
            nodeIds = []
            for blk in qblks:
                if not blk:
                    nodeIds.append(None)
                    continue
                blki, blk_start, blk_end, blk_succ, blk_pred = blk	
                nodeId = (blk_start, blk_end)
                nodeidx = self.AddNode(nodeId)
                nodeIds.append(nodeidx)
                self.nodes[nodeidx] = nodeId
        
            ends = []
            for blk in qblks:
                if not blk:
                    continue
                blki, blk_start, blk_end, blk_succ, blk_pred = blk
                if not blki in outlink_blks:
                    if not blk_succ:
                        ends.append(nodeIds[blki])
                    for succ in blk_succ:
                        print("edge %d -> %d" % (nodeIds[blki], nodeIds[succ]))
                        self.AddEdge(nodeIds[blki], nodeIds[succ])
                else:
                    outlinef = outlink_blks[blki]
                    print ('###### Outline func %x' % outlinef.start_ea)
                    outline_start, outline_ends = add_funcchart(outlinef, funcChain)
                    print('###### Outline func %x End' % outlinef.start_ea)
                    print("edge-outline-call %d -> %d" % (nodeIds[blki], outline_start))
                    self.AddEdge(nodeIds[blki], outline_start)
                    for outline_end in outline_ends:
                        for succ in blk_succ:
                            print("edge-outline-ret %d -> %d" % (outline_end, nodeIds[succ]))
                            self.AddEdge(outline_end, nodeIds[succ])
            return nodeIds[0], ends
        add_funcchart(self.func, ())
        for k,v in self.nodes.items():
            print(k, hex(v[0]), hex(v[1]))
        return True

    def OnGetText(self, node_id):
        #return self[node_id]
        print('%d GetText(%x, %x)' % (node_id, self.nodes[node_id][0], self.nodes[node_id][1]))
        def getdis(s, e):
            t = idaapi.disasm_text_t()
            idaapi.gen_disasm_text(t, s, e, False)
            lines = [(l.at.as_idaplace_t(l.at).ea, l.line) for l in t]
            if not lines:
                print(hex(s), hex(e))
            sea = lines[0][0]
            eea = lines[-1][0]
            hdrlines = 0   
            while hdrlines < 10 and hdrlines < len(lines):
                if 'loc_' in lines[hdrlines][1]:
                    break
                if lines[hdrlines][0] == sea:
                    hdrlines += 1
                else:
                    break
            taillines = 0
            if eea != sea:
                while taillines < 10 and taillines < len(lines):
                    if lines[- taillines-1][0] == eea:
                        taillines += 1
                    else:
                        break
            print(hdrlines, taillines)
            if hdrlines > 1:
                lines = lines[hdrlines - 1:]
            if taillines > 1:
                lines = lines[:-(taillines - 1)]
            return '\n'.join(c[1] for c in lines)
        return getdis(self.nodes[node_id][0], self.nodes[node_id][1])

    def OnPopup(self, widget, popup_handle):
        # graph closer
        actname = "graph_closer:%s" % self.title
        desc = ida_kernwin.action_desc_t(actname, "Close: %s" % self.title, GraphCloser(self))
        ida_kernwin.attach_dynamic_action_to_popup(None, popup_handle, desc)


        # selection printer
        actname = "selection_printer:%s" % self.title
        desc = ida_kernwin.action_desc_t(actname, "Print selection: %s" % self.title, SelectionPrinter(self))
        ida_kernwin.attach_dynamic_action_to_popup(None, popup_handle, desc)


def show_graph():
    f = ida_funcs.get_func(ida_kernwin.get_screen_ea())
    if not f:
        print("Must be in a function")
        return
    g = MyGraph(f, None)
    if g.Show():
        return g
    else:
        return None

g = show_graph()
if g:
    print("Graph created and displayed!")