Example: clean graph_ir IDA

1 files changed, 59 insertions, 183 deletions

diff --git a/example/ida/graph_ir.py b/example/ida/graph_ir.py
index bb06fd0b..678ae2ff 100644
--- a/example/ida/graph_ir.py
+++ b/example/ida/graph_ir.py
@@ -1,40 +1,55 @@
-import sys
 import os
 import tempfile
 
 from idaapi import GraphViewer
 
 from miasm2.core.bin_stream_ida import bin_stream_ida
-from miasm2.core.asmblock import *
+from miasm2.core.asmblock import expr_is_label, AsmLabel, is_int
 from miasm2.expression.simplifications import expr_simp
-from miasm2.expression.expression import *
-from miasm2.analysis.data_analysis import inter_bloc_flow, \
-    intra_bloc_flow_symbexec
 from miasm2.analysis.data_flow import dead_simp
 from miasm2.ir.ir import AssignBlock, IRBlock
 
 from utils import guess_machine, expr2colorstr
 
 
-def color_irbloc(irbloc):
-    o = []
-    lbl = '%s' % irbloc.label
-    lbl = idaapi.COLSTR(lbl, idaapi.SCOLOR_INSN)
-    o.append(lbl)
-    for assignblk in irbloc.irs:
+# Override Miasm asmblock default label naming convention to shrink block size
+# in IDA
+
+def label_init(self, name="", offset=None):
+    self.fixedblocs = False
+    if is_int(name):
+        name = "loc_%X" % (int(name) & 0xFFFFFFFFFFFFFFFF)
+    self.name = name
+    self.attrib = None
+    if offset is None:
+        self.offset = None
+    else:
+        self.offset = int(offset)
+def label_str(self):
+    if isinstance(self.offset, (int, long)):
+        return "%s:0x%x" % (self.name, self.offset)
+    else:
+        return "%s:%s" % (self.name, str(self.offset))
+
+AsmLabel.__init__ = label_init
+AsmLabel.__str__ = label_str
+
+def color_irblock(irblock):
+    out = []
+    lbl = idaapi.COLSTR(str(irblock.label), idaapi.SCOLOR_INSN)
+    out.append(lbl)
+    for assignblk in irblock.irs:
         for dst, src in sorted(assignblk.iteritems()):
             dst_f = expr2colorstr(ir_arch.arch.regs.all_regs_ids, dst)
             src_f = expr2colorstr(ir_arch.arch.regs.all_regs_ids, src)
-            s = idaapi.COLSTR("%s = %s" % (dst_f, src_f), idaapi.SCOLOR_INSN)
-            o.append('    %s' % s)
-        o.append("")
-    o.pop()
-    i = len(irbloc.irs)
-    s = str('    Dst: %s' % irbloc.dst)
-    s = idaapi.COLSTR(s, idaapi.SCOLOR_RPTCMT)
-    o.append(s)
-
-    return "\n".join(o)
+            line = idaapi.COLSTR("%s = %s" % (dst_f, src_f), idaapi.SCOLOR_INSN)
+            out.append('    %s' % line)
+        out.append("")
+    out.pop()
+    dst = str('    Dst: %s' % irblock.dst)
+    dst = idaapi.COLSTR(dst, idaapi.SCOLOR_RPTCMT)
+    out.append(dst)
+    return "\n".join(out)
 
 
 class GraphMiasmIR(GraphViewer):
@@ -50,30 +65,29 @@ class GraphMiasmIR(GraphViewer):
         print 'refresh'
         self.Clear()
         addr_id = {}
-        for irbloc in self.ir_arch.blocks.values():
-            id_irbloc = self.AddNode(color_irbloc(irbloc))
-            addr_id[irbloc] = id_irbloc
+        for irblock in self.ir_arch.blocks.values():
+            id_irblock = self.AddNode(color_irblock(irblock))
+            addr_id[irblock] = id_irblock
 
-        for irbloc in self.ir_arch.blocks.values():
-            if not irbloc:
+        for irblock in self.ir_arch.blocks.values():
+            if not irblock:
                 continue
-            dst = ir_arch.dst_trackback(irbloc)
-            for d in dst:
-                if not expr_is_label(d):
+            all_dst = ir_arch.dst_trackback(irblock)
+            for dst in all_dst:
+                if not expr_is_label(dst):
                     continue
 
-                d = d.name
-                if not d in self.ir_arch.blocks:
+                dst = dst.name
+                if not dst in self.ir_arch.blocks:
                     continue
-                b = self.ir_arch.blocks[d]
-                node1 = addr_id[irbloc]
-                node2 = addr_id[b]
+                dst_block = self.ir_arch.blocks[dst]
+                node1 = addr_id[irblock]
+                node2 = addr_id[dst_block]
                 self.AddEdge(node1, node2)
         return True
 
     def OnGetText(self, node_id):
-        b = self[node_id]
-        return str(b)
+        return str(self[node_id])
 
     def OnSelect(self, node_id):
         return True
@@ -109,31 +123,30 @@ mdis = dis_engine(bs)
 ir_arch = ira(mdis.symbol_pool)
 
 # populate symbols with ida names
-for ad, name in Names():
+for addr, name in Names():
     # print hex(ad), repr(name)
     if name is None:
         continue
-    mdis.symbol_pool.add_label(name, ad)
+    mdis.symbol_pool.add_label(name, addr)
 
 print "start disasm"
-ad = ScreenEA()
-print hex(ad)
+addr = ScreenEA()
+print hex(addr)
 
-ab = mdis.dis_multibloc(ad)
+blocks = mdis.dis_multibloc(addr)
 
 print "generating graph"
-open('asm_flow.dot', 'w').write(ab.dot())
-
+open('asm_flow.dot', 'w').write(blocks.dot())
 
-print "generating IR... %x" % ad
+print "generating IR... %x" % addr
 
-for block in ab:
+for block in blocks:
     print 'ADD'
     print block
     ir_arch.add_bloc(block)
 
 
-print "IR ok... %x" % ad
+print "IR ok... %x" % addr
 
 for irb in ir_arch.blocks.itervalues():
     irs = []
@@ -153,145 +166,8 @@ open(os.path.join(tempfile.gettempdir(), 'graph.dot'), 'wb').write(out)
 
 g = GraphMiasmIR(ir_arch, "Miasm IR graph", None)
 
-
-def mycb(*test):
-    print test
-    raise NotImplementedError('not fully functional')
-
 g.cmd_a = g.AddCommand("cmd a", "x")
 g.cmd_b = g.AddCommand("cmd b", "y")
 
 g.Show()
 
-
-def node2str(n):
-    label, i, node = n
-    print n
-    # out = "%s,%s\n%s"%n
-    out = "%s" % node
-    return out
-
-
-def get_node_name(label, i, n):
-    # n_name = "%s_%d_%s"%(label.name, i, n)
-    n_name = (label.name, i, n)
-    return n_name
-
-
-def get_modified_symbols(sb):
-    # get modified IDS
-    ids = sb.symbols.symbols_id.keys()
-    ids.sort()
-    out = {}
-    for i in ids:
-        if i in sb.arch.regs.regs_init and \
-                i in sb.symbols.symbols_id and \
-                sb.symbols.symbols_id[i] == sb.arch.regs.regs_init[i]:
-            continue
-        # print i, sb.symbols.symbols_id[i]
-        out[i] = sb.symbols.symbols_id[i]
-
-    # get mem IDS
-    mems = sb.symbols.symbols_mem.values()
-    for m, v in mems:
-        # print m, v
-        out[m] = v
-    pp([(str(x[0]), str(x[1])) for x in out.items()])
-    return out
-
-
-def gen_bloc_data_flow_graph(ir_arch, in_str, ad):  # arch, attrib, pool_bin, bloc, symbol_pool):
-    out_str = ""
-
-    # dead_simp(ir_arch)
-
-    irbloc_0 = None
-    for irbloc in ir_arch.blocks.values():
-        if irbloc.label.offset == ad:
-            irbloc_0 = irbloc
-            break
-    assert(irbloc_0 is not None)
-    flow_graph = DiGraph()
-    done = set()
-    todo = set([irbloc_0.label])
-
-    bloc2w = {}
-
-    for irbloc in ir_arch.blocks.values():
-        intra_bloc_flow_symbexec(ir_arch, flow_graph, irbloc)
-        # intra_bloc_flow_symb(ir_arch, flow_graph, irbloc)
-
-    for irbloc in ir_arch.blocks.values():
-        print irbloc
-        print 'IN', [str(x) for x in irbloc.in_nodes]
-        print 'OUT', [str(x) for x in irbloc.out_nodes]
-
-    print '*' * 20, 'interbloc', '*' * 20
-    inter_bloc_flow(ir_arch, flow_graph, irbloc_0.label, False)
-
-    print 'Dataflow roots:'
-    for node in flow_graph.roots():
-        lbl, i, n = node
-        if n in ir_arch.arch.regs.all_regs_ids:
-            print node
-
-    open('data.dot', 'w').write(flow_graph.dot())
-    return flow_graph
-
-
-class GraphMiasmIRFlow(GraphViewer):
-
-    def __init__(self, flow_graph, title, result):
-        GraphViewer.__init__(self, title)
-        print 'init'
-        self.flow_graph = flow_graph
-        self.result = result
-        self.names = {}
-
-    def OnRefresh(self):
-        print 'refresh'
-        self.Clear()
-        addr_id = {}
-        for n in self.flow_graph.nodes():
-            id_n = self.AddNode(node2str(self.flow_graph, n))
-            addr_id[n] = id_n
-
-        for a, b in self.flow_graph.edges():
-                node1, node2 = addr_id[a], addr_id[b]
-                self.AddEdge(node1, node2)
-        return True
-
-    def OnGetText(self, node_id):
-        b = self[node_id]
-        return str(b).lower()
-
-    def OnSelect(self, node_id):
-        return True
-
-    def OnClick(self, node_id):
-        return True
-
-    def OnCommand(self, cmd_id):
-        if self.cmd_test == cmd_id:
-            print 'TEST!'
-            return
-        print "command:", cmd_id
-
-    def Show(self):
-        if not GraphViewer.Show(self):
-            return False
-        self.cmd_test = self.AddCommand("Test", "F2")
-        if self.cmd_test == 0:
-            print "Failed to add popup menu item!"
-        return True
-
-
-#print "gen bloc data flow"
-#flow_graph = gen_bloc_data_flow_graph(ir_arch, bs, ad)
-#def node2str(self, n):
-#    return "%s, %s\\l%s" % n
-#flow_graph.node2str = lambda n: node2str(flow_graph, n)
-#open('data_flow.dot', 'w').write(flow_graph.dot())
-
-# h =  GraphMiasmIRFlow(flow_graph, "Miasm IRFlow graph", None)
-# h.Show()