Example/ida: use addr to guess arch

Some arch like ARM depends on address in order to determine which
attributes have to be used during analysis

5 files changed, 54 insertions, 43 deletions

diff --git a/example/ida/ctype_propagation.py b/example/ida/ctype_propagation.py
index f459022e..d35835dc 100644
--- a/example/ida/ctype_propagation.py
+++ b/example/ida/ctype_propagation.py
@@ -10,7 +10,7 @@ from miasm2.arch.x86.ctype import CTypeAMD64_unk, CTypeX86_unk
 from miasm2.arch.msp430.ctype import CTypeMSP430_unk
 from miasm2.core.objc import CTypesManagerNotPacked, ExprToAccessC, CHandler
 from miasm2.core.ctypesmngr import CAstTypes
-from miasm2.expression.expression import ExprId, ExprInt, ExprOp, ExprAff
+from miasm2.expression.expression import ExprLoc, ExprInt, ExprOp, ExprAff
 from miasm2.ir.symbexec_types import SymbExecCType
 from miasm2.expression.parser import str_to_expr
 from miasm2.analysis.cst_propag import add_state, propagate_cst_expr
@@ -19,9 +19,7 @@ from utils import guess_machine
 
 class TypePropagationForm(ida_kernwin.Form):
 
-    def __init__(self, ira):
-
-        self.ira = ira
+    def __init__(self):
 
         default_types_info = r"""ExprId("RDX", 64): char *"""
         archs = ["AMD64_unk", "X86_32_unk", "msp430_unk"]
@@ -204,7 +202,6 @@ class SymbExecCTypeFix(SymbExecCType):
                     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))
-
             self.eval_updt_assignblk(assignblk)
         for offset, value in offset2cmt.iteritems():
             idc.MakeComm(offset, '\n'.join(value))
@@ -243,38 +240,39 @@ def get_ira_call_fixer(ira):
 
 
 def analyse_function():
-
-    # Init
-    machine = guess_machine()
-    mn, dis_engine, ira = machine.mn, machine.dis_engine, machine.ira
-
-    bs = bin_stream_ida()
-    mdis = dis_engine(bs, dont_dis_nulstart_bloc=True)
-
-
-    iraCallStackFixer = get_ira_call_fixer(ira)
-    ir_arch = iraCallStackFixer(mdis.symbol_pool)
-
-
     # Get settings
-    settings = TypePropagationForm(ir_arch)
+    settings = TypePropagationForm()
     ret = settings.Execute()
     if not ret:
         return
 
+
+    end = None
     if settings.cScope.value == 0:
         addr = settings.functionAddr.value
     else:
         addr = settings.startAddr.value
         if settings.cScope.value == 2:
             end = settings.endAddr
-            mdis.dont_dis = [end]
 
-    blocks = mdis.dis_multiblock(addr)
+    # Init
+    machine = guess_machine(addr=addr)
+    mn, dis_engine, ira = machine.mn, machine.dis_engine, machine.ira
+
+    bs = bin_stream_ida()
+    mdis = dis_engine(bs, dont_dis_nulstart_bloc=True)
+    if end is not None:
+        mdis.dont_dis = [end]
+
+
+    iraCallStackFixer = get_ira_call_fixer(ira)
+    ir_arch = iraCallStackFixer(mdis.symbol_pool)
+
+
+    asmcfg = mdis.dis_multiblock(addr)
     # Generate IR
-    for block in blocks:
+    for block in asmcfg.blocks:
         ir_arch.add_block(block)
-
     cst_propag_link = {}
     if settings.cUnalias.value:
         init_infos = {ir_arch.sp: ir_arch.arch.regs.regs_init[ir_arch.sp] }
@@ -298,7 +296,8 @@ def analyse_function():
         expr_str, ctype_str = expr_str.strip(), ctype_str.strip()
         expr = str_to_expr(expr_str)
         ast = mychandler.types_mngr.types_ast.parse_c_type(
-            ctype_str)
+            ctype_str
+        )
         ctype = mychandler.types_mngr.types_ast.ast_parse_declaration(ast.ext[0])
         objc = types_mngr.get_objc(ctype)
         print '=' * 20
@@ -309,12 +308,15 @@ def analyse_function():
     lbl_real_start = ir_arch.symbol_pool.getby_offset(addr)
     lbl_head = ir_arch.symbol_pool.getby_name_create("start")
 
-    first_block = blocks.label2block(lbl_real_start)
+    first_block = asmcfg.label2block(lbl_real_start)
 
-    assignblk_head = AssignBlock([ExprAff(ir_arch.IRDst, ExprId(lbl_real_start, ir_arch.IRDst.size)),
-                                  ExprAff(
-                                      ir_arch.sp, ir_arch.arch.regs.regs_init[ir_arch.sp])
-                                  ], first_block.lines[0])
+    assignblk_head = AssignBlock(
+        [
+            ExprAff(ir_arch.IRDst, ExprLoc(lbl_real_start, ir_arch.IRDst.size)),
+            ExprAff(ir_arch.sp, ir_arch.arch.regs.regs_init[ir_arch.sp])
+        ],
+        first_block.lines[0]
+    )
     irb_head = IRBlock(lbl_head, [assignblk_head])
     ir_arch.blocks[lbl_head] = irb_head
     ir_arch.graph.add_uniq_edge(lbl_head, lbl_real_start)
@@ -332,7 +334,6 @@ def analyse_function():
         done.add((lbl, state))
         if lbl not in ir_arch.blocks:
             continue
-
         symbexec_engine = TypePropagationEngine(ir_arch, types_mngr, state)
         addr = symbexec_engine.run_block_at(lbl)
         symbexec_engine.del_mem_above_stack(ir_arch.sp)
diff --git a/example/ida/depgraph.py b/example/ida/depgraph.py
index 825d7b90..4320be8d 100644
--- a/example/ida/depgraph.py
+++ b/example/ida/depgraph.py
@@ -198,8 +198,12 @@ def next_element():
 
 def launch_depgraph():
     global graphs, comments, sol_nb, settings, addr, ir_arch
+    # Get the current function
+    addr = idc.ScreenEA()
+    func = ida_funcs.get_func(addr)
+
     # Init
-    machine = guess_machine()
+    machine = guess_machine(addr=func.startEA)
     mn, dis_engine, ira = machine.mn, machine.dis_engine, machine.ira
 
     bs = bin_stream_ida()
@@ -212,9 +216,6 @@ def launch_depgraph():
             continue
         mdis.symbol_pool.add_location(name, ad)
 
-    # Get the current function
-    addr = idc.ScreenEA()
-    func = ida_funcs.get_func(addr)
     asmcfg = mdis.dis_multiblock(func.startEA)
 
     # Generate IR
diff --git a/example/ida/graph_ir.py b/example/ida/graph_ir.py
index 370500e5..97d30851 100644
--- a/example/ida/graph_ir.py
+++ b/example/ida/graph_ir.py
@@ -98,7 +98,9 @@ class GraphMiasmIR(idaapi.GraphViewer):
 
 
 def build_graph(verbose=False, simplify=False):
-    machine = guess_machine()
+    start_addr = idc.ScreenEA()
+
+    machine = guess_machine(addr=start_addr)
     mn, dis_engine, ira = machine.mn, machine.dis_engine, machine.ira
 
     if verbose:
@@ -125,17 +127,16 @@ def build_graph(verbose=False, simplify=False):
 
     if verbose:
         print "start disasm"
-    addr = idc.ScreenEA()
     if verbose:
         print hex(addr)
 
-    asmcfg = mdis.dis_multiblock(addr)
+    asmcfg = mdis.dis_multiblock(start_addr)
 
     if verbose:
         print "generating graph"
         open('asm_flow.dot', 'w').write(asmcfg.dot())
 
-        print "generating IR... %x" % addr
+        print "generating IR... %x" % start_addr
 
     for block in asmcfg.blocks:
         if verbose:
@@ -144,7 +145,7 @@ def build_graph(verbose=False, simplify=False):
         ir_arch.add_block(block)
 
     if verbose:
-        print "IR ok... %x" % addr
+        print "IR ok... %x" % start_addr
 
     for irb in ir_arch.blocks.itervalues():
         irs = []
diff --git a/example/ida/symbol_exec.py b/example/ida/symbol_exec.py
index 49c6fdb6..0d8c63c2 100644
--- a/example/ida/symbol_exec.py
+++ b/example/ida/symbol_exec.py
@@ -128,11 +128,12 @@ def symbolic_exec():
 
     from utils import guess_machine
 
+    start, end = idc.SelStart(), idc.SelEnd()
+
     bs = bin_stream_ida()
-    machine = guess_machine()
+    machine = guess_machine(addr=start)
 
     mdis = machine.dis_engine(bs)
-    start, end = idc.SelStart(), idc.SelEnd()
 
     if start == idc.BADADDR and end == idc.BADADDR:
         start = idc.ScreenEA()
diff --git a/example/ida/utils.py b/example/ida/utils.py
index 481220a9..5620a881 100644
--- a/example/ida/utils.py
+++ b/example/ida/utils.py
@@ -5,7 +5,7 @@ from miasm2.analysis.machine import Machine
 from miasm2.ir.translators import Translator
 import miasm2.expression.expression as m2_expr
 
-def guess_machine():
+def guess_machine(addr=None):
     "Return an instance of Machine corresponding to the IDA guessed processor"
 
     processor_name = GetLongPrm(INF_PROCNAME)
@@ -39,7 +39,14 @@ def guess_machine():
                         (False, 64, True): "aarch64b",
                         (False, 64, False): "aarch64l",
                         }
-        is_armt = globals().get('armt', False)
+
+        # Get T reg to detect arm/thumb function
+        # Default is arm
+        is_armt = False
+        if addr is not None:
+            t_reg = GetReg(addr, "T")
+            is_armt = t_reg == 1
+
         is_bigendian = info.is_be()
         infos = (is_armt, size, is_bigendian)
         if not infos in info2machine: