Merge pull request #537 from commial/fix/ida-examples

Fix/ida examples

2 files changed, 28 insertions, 15 deletions

diff --git a/example/ida/depgraph.py b/example/ida/depgraph.py
index faec2857..cbd0cf0f 100644
--- a/example/ida/depgraph.py
+++ b/example/ida/depgraph.py
@@ -7,6 +7,7 @@ from miasm2.expression import expression as m2_expr
 
 from miasm2.expression.simplifications import expr_simp
 from miasm2.analysis.depgraph import DependencyGraph
+from miasm2.ir.ir import AssignBlock
 
 from utils import guess_machine
 
@@ -21,10 +22,18 @@ class depGraphSettingsForm(Form):
         self.stk_unalias_force = False
 
         self.address = ScreenEA()
-        cur_block = list(ira.getby_offset(self.address))[0]
-        for line_nb, l in enumerate(cur_block.lines):
-            if l.offset == self.address:
+        cur_block = None
+        for block in ira.getby_offset(self.address):
+            if block.label.offset is not None:
+                # Only one block non-generated
+                assert cur_block is None
+                cur_block = block
+        assert cur_block is not None
+        line_nb = None
+        for line_nb, assignblk in enumerate(cur_block.irs):
+            if assignblk.instr.offset == self.address:
                 break
+        assert line_nb is not None
         cur_label = str(cur_block.label)
         labels = sorted(map(str, ira.blocks.keys()))
         regs = sorted(ir_arch.arch.regs.all_regs_ids_byname.keys())
@@ -102,7 +111,7 @@ Method to use:
     def elements(self):
         value = self.cbReg.value
         if value in self.stk_args:
-            line = self.ira.blocks[self.label].lines[self.line_nb]
+            line = self.ira.blocks[self.label].irs[self.line_nb].instr
             arg_num = self.stk_args[value]
             stk_high = m2_expr.ExprInt(GetSpd(line.offset), ir_arch.sp.size)
             stk_off = m2_expr.ExprInt(self.ira.sp.size/8 * arg_num, ir_arch.sp.size)
@@ -167,17 +176,18 @@ for irb in ir_arch.blocks.values():
     fix_stack = irb.label.offset is not None and settings.unalias_stack
     for i, assignblk in enumerate(irb.irs):
         if fix_stack:
-            stk_high = m2_expr.ExprInt(GetSpd(irb.lines[i].offset), ir_arch.sp.size)
+            stk_high = m2_expr.ExprInt(GetSpd(irb.irs[i].instr.offset), ir_arch.sp.size)
             fix_dct = {ir_arch.sp: mn.regs.regs_init[ir_arch.sp] + stk_high}
 
-        for dst, src in assignblk.items():
-            del(assignblk[dst])
+        new_assignblk = {}
+        for dst, src in assignblk.iteritems():
             if fix_stack:
                 src = src.replace_expr(fix_dct)
                 if dst != ir_arch.sp:
                     dst = dst.replace_expr(fix_dct)
             dst, src = expr_simp(dst), expr_simp(src)
-            assignblk[dst] = src
+            new_assignblk[dst] = src
+        irb.irs[i] = AssignBlock(new_assignblk, instr=assignblk.instr)
 
 # Get dependency graphs
 dg = settings.depgraph
@@ -215,7 +225,7 @@ def treat_element():
 
     for node in graph.relevant_nodes:
         try:
-            offset = ir_arch.blocks[node.label].lines[node.line_nb].offset
+            offset = ir_arch.blocks[node.label].irs[node.line_nb].instr.offset
         except IndexError:
             print "Unable to highlight %s" % node
             continue
diff --git a/example/ida/graph_ir.py b/example/ida/graph_ir.py
index d2552c82..6ff4304a 100644
--- a/example/ida/graph_ir.py
+++ b/example/ida/graph_ir.py
@@ -11,6 +11,8 @@ 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
+
 from utils import guess_machine, expr2colorstr
 
 
@@ -133,12 +135,13 @@ for block in ab:
 
 print "IR ok... %x" % ad
 
-for irb in ir_arch.blocks.values():
-    for assignblk in irb.irs:
-        for dst, src in assignblk.items():
-            del(assignblk[dst])
-            dst, src = expr_simp(dst), expr_simp(src)
-            assignblk[dst] = src
+for irb in ir_arch.blocks.itervalues():
+    for i, assignblk in enumerate(irb.irs):
+        new_assignblk = {
+            expr_simp(dst): expr_simp(src)
+            for dst, src in assignblk.iteritems()
+        }
+        irb.irs[i] = AssignBlock(new_assignblk, instr=assignblk.instr)
 
 out = ir_arch.graph.dot()
 open(os.path.join(tempfile.gettempdir(), 'graph.dot'), 'wb').write(out)