Examples/IDA: Export common IDA methods to a file `utils`

1 files changed, 113 insertions, 0 deletions

diff --git a/example/ida/utils.py b/example/ida/utils.py
new file mode 100644
index 00000000..6b42cd83
--- /dev/null
+++ b/example/ida/utils.py
@@ -0,0 +1,113 @@
+import idaapi
+from idc import *
+
+from miasm2.analysis.machine import Machine
+import miasm2.expression.expression as m2_expr
+
+
+def guess_machine():
+    "Return an instance of Machine corresponding to the IDA guessed processor"
+
+    processor_name = GetLongPrm(INF_PROCNAME)
+
+    if processor_name == "metapc":
+
+        # HACK: check 32/64 using INF_START_SP
+        max_size = GetLongPrm(INF_START_SP)
+        if max_size == 0x80:  # TODO XXX check
+            machine = Machine("x86_16")
+        elif max_size == 0xFFFFFFFF:
+            machine = Machine("x86_32")
+        elif max_size == 0xFFFFFFFFFFFFFFFF:
+            machine = Machine("x86_64")
+        else:
+            raise ValueError('cannot guess 32/64 bit! (%x)' % max_size)
+    elif processor_name == "ARM":
+        # TODO ARM/thumb
+        # hack for thumb: set armt = True in globals :/
+        # set bigendiant = True is bigendian
+        is_armt = globals().get('armt', False)
+        is_bigendian = globals().get('bigendian', False)
+        if is_armt:
+            if is_bigendian:
+                machine = Machine("armtb")
+            else:
+                machine = Machine("armtl")
+        else:
+            if is_bigendian:
+                machine = Machine("armb")
+            else:
+                machine = Machine("arml")
+
+        from miasm2.analysis.disasm_cb import arm_guess_subcall, arm_guess_jump_table
+        guess_funcs.append(arm_guess_subcall)
+        guess_funcs.append(arm_guess_jump_table)
+
+    elif processor_name == "msp430":
+        machine = Machine("msp430")
+    elif processor_name == "mipsl":
+        machine = Machine("mipsl")
+    elif processor_name == "mipsb":
+        machine = Machine("mipsb")
+    else:
+        print repr(processor_name)
+        raise NotImplementedError('not fully functional')
+
+    return machine
+
+
+def expr2colorstr(regs_ids, expr):
+    """Colorize an Expr instance for IDA
+    @regs_ids: list of ExprId corresponding to available registers
+    @expr: Expr instance to colorize
+    """
+
+    if isinstance(expr, m2_expr.ExprId):
+        s = str(expr)
+        if expr in regs_ids:
+            s = idaapi.COLSTR(s, idaapi.SCOLOR_REG)
+    elif isinstance(expr, m2_expr.ExprInt):
+        s = str(expr)
+        s = idaapi.COLSTR(s, idaapi.SCOLOR_NUMBER)
+    elif isinstance(expr, m2_expr.ExprMem):
+        s = '%s[%s]' % (idaapi.COLSTR('@' + str(expr.size),
+                                      idaapi.SCOLOR_RPTCMT),
+                         expr2colorstr(regs_ids, expr.arg))
+    elif isinstance(expr, m2_expr.ExprOp):
+        out = []
+        for a in expr.args:
+            s = expr2colorstr(regs_ids, a)
+            if isinstance(a, m2_expr.ExprOp):
+                s = "(%s)" % s
+            out.append(s)
+        if len(out) == 1:
+            s = "%s %s" % (expr.op, str(out[0]))
+        else:
+            s = (" " + expr.op + " ").join(out)
+    elif isinstance(expr, m2_expr.ExprAff):
+        s = "%s = %s" % (
+            expr2colorstr(regs_ids, expr.dst), expr2colorstr(regs_ids, expr.src))
+    elif isinstance(expr, m2_expr.ExprCond):
+        cond = expr2colorstr(regs_ids, expr.cond)
+        src1 = expr2colorstr(regs_ids, expr.src1)
+        src2 = expr2colorstr(regs_ids, expr.src2)
+        s = "(%s?%s:%s)" % (cond, src1, src2)
+    elif isinstance(expr, m2_expr.ExprSlice):
+        s = "(%s)[%s:%s]" % (expr2colorstr(regs_ids, expr.arg),
+                             idaapi.COLSTR(str(expr.start),
+                                           idaapi.SCOLOR_RPTCMT),
+                             idaapi.COLSTR(str(expr.stop),
+                                           idaapi.SCOLOR_RPTCMT))
+    elif isinstance(expr, m2_expr.ExprCompose):
+        s = "{"
+        s += ", ".join(["%s, %s, %s" % (expr2colorstr(regs_ids, subexpr),
+                                        idaapi.COLSTR(str(start),
+                                                      idaapi.SCOLOR_RPTCMT),
+                                        idaapi.COLSTR(str(stop),
+                                                      idaapi.SCOLOR_RPTCMT))
+                        for subexpr, start, stop in expr.args])
+        s += "}"
+    else:
+        s = str(expr)
+
+    return s