about summary refs log tree commit diff stats
path: root/example/ida/utils.py
diff options
context:
space:
mode:
Diffstat (limited to 'example/ida/utils.py')
-rw-r--r--example/ida/utils.py45
1 files changed, 31 insertions, 14 deletions
diff --git a/example/ida/utils.py b/example/ida/utils.py
index e026f2fc..c66475f2 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:
@@ -72,22 +79,29 @@ class TranslatorIDA(Translator):
     # Implemented language
     __LANG__ = "ida_w_color"
 
-    def __init__(self, regs_ids=None, **kwargs):
+    def __init__(self, loc_db=None, **kwargs):
         super(TranslatorIDA, self).__init__(**kwargs)
-        if regs_ids is None:
-            regs_ids = {}
-        self.regs_ids = regs_ids
+        self.loc_db = loc_db
 
     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.loc_db is not None:
+            out = self.loc_db.pretty_str(expr.loc_key)
+        else:
+            out = str(expr)
+        out = idaapi.COLSTR(out, idaapi.SCOLOR_REG)
         return out
 
     def from_ExprMem(self, expr):
@@ -126,20 +140,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, loc_db):
     """Colorize an Expr instance for IDA
-    @regs_ids: list of ExprId corresponding to available registers
     @expr: Expr instance to colorize
+    @loc_db: LocationDB instance
     """
 
-    translator = TranslatorIDA(regs_ids)
+    translator = TranslatorIDA(loc_db=loc_db)
     return translator.from_expr(expr)