about summary refs log tree commit diff stats
path: root/example/ida
diff options
context:
space:
mode:
Diffstat (limited to 'example/ida')
-rw-r--r--example/ida/ctype_propagation.py2
-rw-r--r--example/ida/depgraph.py24
-rw-r--r--example/ida/graph_ir.py22
-rw-r--r--example/ida/symbol_exec.py23
-rw-r--r--example/ida/utils.py34
5 files changed, 62 insertions, 43 deletions
diff --git a/example/ida/ctype_propagation.py b/example/ida/ctype_propagation.py
index 9b9c2e95..f459022e 100644
--- a/example/ida/ctype_propagation.py
+++ b/example/ida/ctype_propagation.py
@@ -201,7 +201,7 @@ class SymbExecCTypeFix(SymbExecCType):
                 if expr.is_int():
                     continue
                 for c_str, c_type in self.chandler.expr_to_c_and_types(expr, self.symbols):
-                    expr = self.cst_propag_link.get((irb.label, index), {}).get(expr, expr)
+                    expr = self.cst_propag_link.get((irb.loc_key, index), {}).get(expr, expr)
                     offset2cmt.setdefault(instr.offset, set()).add(
                         "\n%s: %s\n%s" % (expr, c_str, c_type))
 
diff --git a/example/ida/depgraph.py b/example/ida/depgraph.py
index 5342313a..ece02ad4 100644
--- a/example/ida/depgraph.py
+++ b/example/ida/depgraph.py
@@ -28,7 +28,8 @@ class depGraphSettingsForm(ida_kernwin.Form):
         self.address = idc.ScreenEA()
         cur_block = None
         for block in ira.getby_offset(self.address):
-            if block.label.offset is not None:
+            offset = self.ira.symbol_pool.loc_key_to_offset(block.loc_key)
+            if offset is not None:
                 # Only one block non-generated
                 assert cur_block is None
                 cur_block = block
@@ -38,7 +39,7 @@ class depGraphSettingsForm(ida_kernwin.Form):
             if assignblk.instr.offset == self.address:
                 break
         assert line_nb is not None
-        cur_label = str(cur_block.label)
+        cur_label = str(cur_block.loc_key)
         labels = sorted(map(str, ira.blocks.keys()))
         regs = sorted(ira.arch.regs.all_regs_ids_byname.keys())
         regs += self.stk_args.keys()
@@ -110,13 +111,13 @@ Method to use:
         elif mode == 1:
             return value + 1
         else:
-            return len(self.ira.blocks[self.label])
+            return len(self.ira.blocks[self.loc_key])
 
     @property
     def elements(self):
         value = self.cbReg.value
         if value in self.stk_args:
-            line = self.ira.blocks[self.label][self.line_nb].instr
+            line = self.ira.blocks[self.loc_key][self.line_nb].instr
             arg_num = self.stk_args[value]
             stk_high = m2_expr.ExprInt(idc.GetSpd(line.offset), ir_arch.sp.size)
             stk_off = m2_expr.ExprInt(self.ira.sp.size/8 * arg_num, ir_arch.sp.size)
@@ -174,7 +175,7 @@ def treat_element():
 
     for node in graph.relevant_nodes:
         try:
-            offset = ir_arch.blocks[node.label][node.line_nb].instr.offset
+            offset = ir_arch.blocks[node.loc_key][node.line_nb].instr.offset
         except IndexError:
             print "Unable to highlight %s" % node
             continue
@@ -209,26 +210,27 @@ def launch_depgraph():
     for ad, name in idautils.Names():
         if name is None:
             continue
-        mdis.symbol_pool.add_label(name, ad)
+        mdis.symbol_pool.add_location(name, ad)
 
     # Get the current function
     addr = idc.ScreenEA()
     func = ida_funcs.get_func(addr)
-    blocks = mdis.dis_multiblock(func.startEA)
+    asmcfg = mdis.dis_multiblock(func.startEA)
 
     # Generate IR
-    for block in blocks:
+    for block in asmcfg.blocks:
         ir_arch.add_block(block)
 
     # Get settings
     settings = depGraphSettingsForm(ir_arch)
     settings.Execute()
 
-    label, elements, line_nb = settings.label, settings.elements, settings.line_nb
+    label, elements, line_nb = settings.loc_key, settings.elements, settings.line_nb
     # Simplify affectations
     for irb in ir_arch.blocks.values():
         irs = []
-        fix_stack = irb.label.offset is not None and settings.unalias_stack
+        offset = ir_arch.symbol_pool.loc_key_to_offset(irb.loc_key)
+        fix_stack = offset is not None and settings.unalias_stack
         for assignblk in irb:
             if fix_stack:
                 stk_high = m2_expr.ExprInt(idc.GetSpd(assignblk.instr.offset), ir_arch.sp.size)
@@ -243,7 +245,7 @@ def launch_depgraph():
                 dst, src = expr_simp(dst), expr_simp(src)
                 new_assignblk[dst] = src
             irs.append(AssignBlock(new_assignblk, instr=assignblk.instr))
-        ir_arch.blocks[irb.label] = IRBlock(irb.label, irs)
+        ir_arch.blocks[irb.loc_key] = IRBlock(irb.loc_key, irs)
 
     # Get dependency graphs
     dg = settings.depgraph
diff --git a/example/ida/graph_ir.py b/example/ida/graph_ir.py
index fad793ff..370500e5 100644
--- a/example/ida/graph_ir.py
+++ b/example/ida/graph_ir.py
@@ -6,7 +6,7 @@ import idc
 import idautils
 
 from miasm2.core.bin_stream_ida import bin_stream_ida
-from miasm2.core.asmblock import AsmLabel, is_int
+from miasm2.core.asmblock import is_int
 from miasm2.expression.simplifications import expr_simp
 from miasm2.analysis.data_flow import dead_simp
 from miasm2.ir.ir import AssignBlock, IRBlock
@@ -33,17 +33,15 @@ def label_str(self):
     else:
         return "%s:%s" % (self.name, str(self.offset))
 
-AsmLabel.__init__ = label_init
-AsmLabel.__str__ = label_str
 
 def color_irblock(irblock, ir_arch):
     out = []
-    lbl = idaapi.COLSTR(str(irblock.label), idaapi.SCOLOR_INSN)
+    lbl = idaapi.COLSTR(ir_arch.symbol_pool.str_loc_key(irblock.loc_key), idaapi.SCOLOR_INSN)
     out.append(lbl)
     for assignblk in irblock:
         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)
+            dst_f = expr2colorstr(dst, symbol_pool=ir_arch.symbol_pool)
+            src_f = expr2colorstr(src, symbol_pool=ir_arch.symbol_pool)
             line = idaapi.COLSTR("%s = %s" % (dst_f, src_f), idaapi.SCOLOR_INSN)
             out.append('    %s' % line)
         out.append("")
@@ -74,7 +72,7 @@ class GraphMiasmIR(idaapi.GraphViewer):
                 continue
             all_dst = self.ir_arch.dst_trackback(irblock)
             for dst in all_dst:
-                if not dst.is_label():
+                if not dst.is_loc():
                     continue
                 if not dst.loc_key in self.ir_arch.blocks:
                     continue
@@ -123,7 +121,7 @@ def build_graph(verbose=False, simplify=False):
             mdis.symbol_pool.getby_name(name)):
             # Symbol alias
             continue
-        mdis.symbol_pool.add_label(name, addr)
+        mdis.symbol_pool.add_location(name, addr)
 
     if verbose:
         print "start disasm"
@@ -131,15 +129,15 @@ def build_graph(verbose=False, simplify=False):
     if verbose:
         print hex(addr)
 
-    blocks = mdis.dis_multiblock(addr)
+    asmcfg = mdis.dis_multiblock(addr)
 
     if verbose:
         print "generating graph"
-        open('asm_flow.dot', 'w').write(blocks.dot())
+        open('asm_flow.dot', 'w').write(asmcfg.dot())
 
         print "generating IR... %x" % addr
 
-    for block in blocks:
+    for block in asmcfg.blocks:
         if verbose:
             print 'ADD'
             print block
@@ -156,7 +154,7 @@ def build_graph(verbose=False, simplify=False):
                 for dst, src in assignblk.iteritems()
             }
             irs.append(AssignBlock(new_assignblk, instr=assignblk.instr))
-        ir_arch.blocks[irb.label] = IRBlock(irb.label, irs)
+        ir_arch.blocks[irb.loc_key] = IRBlock(irb.loc_key, irs)
 
     if verbose:
         out = ir_arch.graph.dot()
diff --git a/example/ida/symbol_exec.py b/example/ida/symbol_exec.py
index f019f77d..63014ece 100644
--- a/example/ida/symbol_exec.py
+++ b/example/ida/symbol_exec.py
@@ -34,8 +34,16 @@ class ActionHandlerTranslate(ActionHandler):
 class symbolicexec_t(idaapi.simplecustviewer_t):
 
     def add(self, key, value):
-        self.AddLine("%s = %s" % (expr2colorstr(self.machine.mn.regs.all_regs_ids, key),
-                                  expr2colorstr(self.machine.mn.regs.all_regs_ids, value)))
+        self.AddLine("%s = %s" % (
+            expr2colorstr(
+                key,
+                symbol_pool=self.symbol_pool
+            ),
+            expr2colorstr(
+                value,
+                symbol_pool=self.symbol_pool
+            )
+        ))
 
     def expand(self, linenum):
         element = self.line2eq[linenum]
@@ -61,11 +69,12 @@ class symbolicexec_t(idaapi.simplecustviewer_t):
         form.Compile()
         form.Execute()
 
-    def Create(self, equations, machine, *args, **kwargs):
+    def Create(self, equations, machine, symbol_pool, *args, **kwargs):
         if not super(symbolicexec_t, self).Create(*args, **kwargs):
             return False
 
         self.machine = machine
+        self.symbol_pool = symbol_pool
         self.line2eq = sorted(equations.items(), key=operator.itemgetter(0))
         self.lines_expanded = set()
 
@@ -126,9 +135,9 @@ def symbolic_exec():
     start, end = idc.SelStart(), idc.SelEnd()
 
     mdis.dont_dis = [end]
-    blocks = mdis.dis_multiblock(start)
-    ira = machine.ira()
-    for block in blocks:
+    asmcfg = mdis.dis_multiblock(start)
+    ira = machine.ira(symbol_pool=mdis.symbol_pool)
+    for block in asmcfg.blocks:
         ira.add_block(block)
 
     print "Run symbolic execution..."
@@ -141,7 +150,7 @@ def symbolic_exec():
 
     view = symbolicexec_t()
     all_views.append(view)
-    if not view.Create(modified, machine,
+    if not view.Create(modified, machine, mdis.symbol_pool,
                        "Symbolic Execution - 0x%x to 0x%x" % (start, end)):
         return
 
diff --git a/example/ida/utils.py b/example/ida/utils.py
index e026f2fc..481220a9 100644
--- a/example/ida/utils.py
+++ b/example/ida/utils.py
@@ -72,22 +72,29 @@ class TranslatorIDA(Translator):
     # Implemented language
     __LANG__ = "ida_w_color"
 
-    def __init__(self, regs_ids=None, **kwargs):
+    def __init__(self, symbol_pool=None, **kwargs):
         super(TranslatorIDA, self).__init__(**kwargs)
-        if regs_ids is None:
-            regs_ids = {}
-        self.regs_ids = regs_ids
+        self.symbol_pool = symbol_pool
 
     def str_protected_child(self, child, parent):
-        return ("(%s)" % self.from_expr(child)) if m2_expr.should_parenthesize_child(child, parent) else self.from_expr(child)
+        return ("(%s)" % (
+            self.from_expr(child)) if m2_expr.should_parenthesize_child(child, parent)
+                else self.from_expr(child)
+        )
 
     def from_ExprInt(self, expr):
         return idaapi.COLSTR(str(expr), idaapi.SCOLOR_NUMBER)
 
     def from_ExprId(self, expr):
-        out = str(expr)
-        if expr in self.regs_ids:
-            out = idaapi.COLSTR(out, idaapi.SCOLOR_REG)
+        out = idaapi.COLSTR(str(expr), idaapi.SCOLOR_REG)
+        return out
+
+    def from_ExprLoc(self, expr):
+        if self.symbol_pool is not None:
+            out = self.symbol_pool.str_loc_key(expr.loc_key)
+        else:
+            out = str(expr)
+        out = idaapi.COLSTR(out, idaapi.SCOLOR_REG)
         return out
 
     def from_ExprMem(self, expr):
@@ -126,20 +133,23 @@ class TranslatorIDA(Translator):
             return (' ' + expr._op + ' ').join([self.str_protected_child(arg, expr)
                                                 for arg in expr._args])
         return (expr._op + '(' +
-                ', '.join([self.from_expr(arg) for arg in expr._args]) + ')')
+                ', '.join(
+                    self.from_expr(arg)
+                    for arg in expr._args
+                ) + ')')
 
     def from_ExprAff(self, expr):
         return "%s = %s" % tuple(map(expr.from_expr, (expr.dst, expr.src)))
 
 
 
-def expr2colorstr(regs_ids, expr):
+def expr2colorstr(expr, symbol_pool):
     """Colorize an Expr instance for IDA
-    @regs_ids: list of ExprId corresponding to available registers
     @expr: Expr instance to colorize
+    @symbol_pool: AsmSymbolPool instance
     """
 
-    translator = TranslatorIDA(regs_ids)
+    translator = TranslatorIDA(symbol_pool=symbol_pool)
     return translator.from_expr(expr)