Merge pull request #52 from commial/ida-symbol_exec

Ida symbol exec

5 files changed, 360 insertions, 96 deletions

diff --git a/example/ida/graph_ir.py b/example/ida/graph_ir.py
index 5f9aead6..a306eda5 100644
--- a/example/ida/graph_ir.py
+++ b/example/ida/graph_ir.py
@@ -16,44 +16,7 @@ from miasm2.analysis.machine import Machine
 from miasm2.analysis.data_analysis import intra_bloc_flow_raw, inter_bloc_flow
 from miasm2.analysis.data_analysis import intra_bloc_flow_symbexec
 
-
-
-def expr2colorstr(ir_arch, e):
-    # print "XXX", e
-    if isinstance(e, ExprId):
-        s = str(e)
-        if e in ir_arch.arch.regs.all_regs_ids:
-            s = idaapi.COLSTR(s, idaapi.SCOLOR_REG)
-    elif isinstance(e, ExprInt):
-        s = str(e)
-        s = idaapi.COLSTR(s, idaapi.SCOLOR_NUMBER)
-    elif isinstance(e, ExprMem):
-        s = '@%d[%s]' % (e.size, expr2colorstr(ir_arch, e.arg))
-    elif isinstance(e, ExprOp):
-        out = []
-        for a in e.args:
-            s = expr2colorstr(ir_arch, a)
-            if isinstance(a, ExprOp):
-                s = "(%s)" % s
-            out.append(s)
-        if len(out) == 1:
-            s = "%s %s" % (e.op, str(out[0]))
-        else:
-            s = (" " + e.op + " ").join(out)
-    elif isinstance(e, ExprAff):
-        s = "%s = %s" % (
-            expr2colorstr(ir_arch, e.dst), expr2colorstr(ir_arch, e.src))
-    elif isinstance(e, ExprCond):
-        cond = expr2colorstr(ir_arch, e.cond)
-        src1 = expr2colorstr(ir_arch, e.src1)
-        src2 = expr2colorstr(ir_arch, e.src2)
-        s = "(%s?%s:%s)" % (cond, src1, src2)
-    elif isinstance(e, ExprSlice):
-        s = "(%s)[%d:%d]" % (expr2colorstr(ir_arch, e.arg), e.start, e.stop)
-    else:
-        s = str(e)
-    # print repr(s)
-    return s
+from utils import guess_machine, expr2colorstr
 
 
 def color_irbloc(irbloc):
@@ -63,7 +26,7 @@ def color_irbloc(irbloc):
     o.append(lbl)
     for i, expr in enumerate(irbloc.irs):
         for e in expr:
-            s = expr2colorstr(ir_arch, e)
+            s = expr2colorstr(ir_arch.arch.regs.all_regs_ids, e)
             s = idaapi.COLSTR(s, idaapi.SCOLOR_INSN)
             o.append('    %s' % s)
         o.append("")
@@ -137,52 +100,7 @@ class GraphMiasmIR(GraphViewer):
 
 from miasm2.analysis.disasm_cb import guess_funcs, guess_multi_cb
 
-
-processor_name = GetLongPrm(INF_PROCNAME)
-dis_engine = None
-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')
-
+machine = guess_machine()
 mn, dis_engine, ira = machine.mn, machine.dis_engine, machine.ira
 
 print "Arch", dis_engine
diff --git a/example/ida/symbol_exec.py b/example/ida/symbol_exec.py
new file mode 100644
index 00000000..c7aff5b5
--- /dev/null
+++ b/example/ida/symbol_exec.py
@@ -0,0 +1,125 @@
+import operator
+
+import idaapi
+import idc
+from miasm2.expression.expression_helper import Variables_Identifier
+from miasm2.expression.expression import ExprAff
+
+from utils import expr2colorstr, translatorForm
+
+
+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)))
+
+    def expand(self, linenum):
+        element = self.line2eq[linenum]
+        expanded = Variables_Identifier(element[1],
+                                        var_prefix="%s_v" % element[0])
+        self.line2eq = self.line2eq[0:linenum] + \
+            expanded.vars.items() + \
+            [(element[0], expanded.equation)] + \
+            self.line2eq[linenum + 1:]
+
+    def print_lines(self):
+        self.ClearLines()
+
+        for element in self.line2eq:
+            self.add(*element)
+
+        self.Refresh()
+
+    def translate_expr(self, line_nb):
+        element = self.line2eq[line_nb]
+        expr = ExprAff(*element)
+        form = translatorForm(expr)
+        form.Compile()
+        form.Execute()
+
+    def Create(self, equations, machine, *args, **kwargs):
+        if not super(symbolicexec_t, self).Create(*args, **kwargs):
+            return False
+
+        self.machine = machine
+        self.line2eq = sorted(equations.items(), key=operator.itemgetter(0))
+        self.lines_expanded = set()
+
+        self.print_lines()
+
+        self.menu_expand = self.AddPopupMenu("Expand [E]")
+        self.menu_translate = self.AddPopupMenu("Translate [T]")
+        return True
+
+    def OnPopupMenu(self, menu_id):
+        if menu_id == self.menu_expand:
+            self.expand(self.GetLineNo())
+            self.print_lines()
+        if menu_id == self.menu_translate:
+            self.translate_expr(self.GetLineNo())
+        return True
+
+    def OnKeydown(self, vkey, shift):
+        # ESCAPE
+        if vkey == 27:
+            self.Close()
+            return True
+        # E (expand)
+        if vkey == 69:
+            self.OnPopupMenu(self.menu_expand)
+        # T (translate)
+        if vkey == 84:
+            self.OnPopupMenu(self.menu_translate)
+        return False
+
+
+def symbolic_exec():
+    from miasm2.analysis.machine import Machine
+    from miasm2.ir.symbexec import symbexec
+    from miasm2.core.bin_stream_ida import bin_stream_ida
+
+    from utils import guess_machine
+
+    bs = bin_stream_ida()
+    machine = guess_machine()
+
+    mdis = machine.dis_engine(bs)
+    start, end = SelStart(), SelEnd()
+
+    mdis.dont_dis = [end]
+    blocs = mdis.dis_multibloc(start)
+    ira = machine.ira()
+    for bloc in blocs:
+        ira.add_bloc(bloc)
+
+    print "Run symbolic execution..."
+    sb = symbexec(ira, machine.mn.regs.regs_init)
+    sb.emul_ir_blocs(ira, start)
+
+    modified = {}
+    for ident in sb.symbols.symbols_id:
+        if ident in sb.ir_arch.arch.regs.regs_init and \
+                ident in sb.symbols.symbols_id and \
+                sb.symbols.symbols_id[ident] == sb.ir_arch.arch.regs.regs_init[ident]:
+            continue
+        modified[ident] = sb.symbols.symbols_id[ident]
+
+    for ident in sb.symbols.symbols_mem:
+        modified[sb.symbols.symbols_mem[ident][0]] = sb.symbols.symbols_mem[ident][1]
+
+
+    view = symbolicexec_t()
+    if not view.Create(modified, machine,
+                       "Symbolic Execution - 0x%x to 0x%x" % (start, end)):
+        return
+
+    view.Show()
+
+idaapi.CompileLine('static key_F3() { RunPythonStatement("symbolic_exec()"); }')
+idc.AddHotkey("F3", "key_F3")
+
+print "=" * 50
+print """Available commands:
+    symbolic_exec() - F3: Symbolic execution of current selection
+"""
diff --git a/example/ida/utils.py b/example/ida/utils.py
new file mode 100644
index 00000000..8b9dcf6a
--- /dev/null
+++ b/example/ida/utils.py
@@ -0,0 +1,169 @@
+import idaapi
+from idc import *
+
+from miasm2.analysis.machine import Machine
+from miasm2.ir.translators import Translator
+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
+
+
+class translatorForm(idaapi.Form):
+    """Translator Form.
+
+    Offer a ComboBox with available languages (ie. IR translators) and the
+    corresponding translation."""
+
+    flags = (idaapi.Form.MultiLineTextControl.TXTF_FIXEDFONT | \
+                 idaapi.Form.MultiLineTextControl.TXTF_READONLY)
+
+    def __init__(self, expr):
+        "@expr: Expr instance"
+
+        # Init
+        self.languages = list(Translator.available_languages())
+        self.expr = expr
+
+        # Initial translation
+        text = Translator.to_language(self.languages[0]).from_expr(self.expr)
+
+        # Create the Form
+        idaapi.Form.__init__(self, r"""STARTITEM 0
+Python Expression
+{FormChangeCb}
+<Language:{cbLanguage}>
+<Translation:{result}>
+""", {
+            'result': idaapi.Form.MultiLineTextControl(text=text,
+                                                       flags=translatorForm.flags),
+            'cbLanguage': idaapi.Form.DropdownListControl(
+                    items=self.languages,
+                    readonly=True,
+                    selval=0),
+            'FormChangeCb': idaapi.Form.FormChangeCb(self.OnFormChange),
+        })
+
+    def OnFormChange(self, fid):
+        if fid == self.cbLanguage.id:
+            # Display the Field (may be hide)
+            self.ShowField(self.result, True)
+
+            # Translate the expression
+            dest_lang = self.languages[self.GetControlValue(self.cbLanguage)]
+            try:
+                text = Translator.to_language(dest_lang).from_expr(self.expr)
+            except Exception, error:
+                self.ShowField(self.result, False)
+                return -1
+
+            # Update the form
+            self.SetControlValue(self.result,
+                                 idaapi.textctrl_info_t(text=str(text),
+                                                        flags=translatorForm.flags))
+        return 1
diff --git a/miasm2/expression/expression_helper.py b/miasm2/expression/expression_helper.py
index 2f0bd4e7..825cad60 100644
--- a/miasm2/expression/expression_helper.py
+++ b/miasm2/expression/expression_helper.py
@@ -17,7 +17,6 @@
 #
 
 # Expressions manipulation functions
-import re
 import itertools
 import collections
 import random
@@ -210,16 +209,20 @@ class Variables_Identifier(object):
     - original expression with variables translated
     """
 
-    var_identifier = re.compile("v\d+")
+    # Attribute used to distinguish created variables from original ones
+    is_var_ident = "is_var_ident"
 
-    def __init__(self, expr):
+    def __init__(self, expr, var_prefix="v"):
         """Set the expression @expr to handle and launch variable identification
-        process"""
+        process
+        @expr: Expr instance
+        @var_prefix: (optional) prefix of the variable name, default is 'v'"""
 
         # Init
         self.var_indice = itertools.count()
         self.var_asked = set()
         self._vars = {} # VarID -> Expr
+        self.var_prefix = var_prefix
 
         # Launch recurrence
         self.find_variables_rec(expr)
@@ -254,9 +257,13 @@ class Variables_Identifier(object):
 
         ## Build initial needs
         for var_id, var_expr in self._vars.iteritems():
+            ### Handle corner cases while using Variable Identifier on an
+            ### already computed equation
             needs[var_id] = [var_name
                              for var_name in var_expr.get_r(mem_read=True)
-                             if self.is_var_identifier(var_name)]
+                             if self.is_var_identifier(var_name) and \
+                                 var_name in todo and \
+                                 var_name != var_id]
 
         ## Build order list
         while todo:
@@ -268,7 +275,6 @@ class Variables_Identifier(object):
                         # A dependency is not met
                         all_met = False
                         break
-
                 if not all_met:
                     continue
 
@@ -282,12 +288,12 @@ class Variables_Identifier(object):
 
     @classmethod
     def is_var_identifier(cls, expr):
-        "Return True iff expr seems to be a variable identifier"
+        "Return True iff @expr is a variable identifier"
         if not isinstance(expr, m2_expr.ExprId):
             return False
 
-        match = cls.var_identifier.match(expr.name)
-        return match is not None and match.group(0) == expr.name
+        return hasattr(expr, cls.is_var_ident) and \
+            getattr(expr, cls.is_var_ident) == True
 
     def find_variables_rec(self, expr):
         """Recursive method called by find_variable to expand @expr.
@@ -301,8 +307,10 @@ class Variables_Identifier(object):
 
             if (expr not in self._vars.values()):
                 # Create var
-                identifier = m2_expr.ExprId("v%s" % self.var_indice.next(),
-                                    size = expr.size)
+                identifier = m2_expr.ExprId("%s%s" % (self.var_prefix,
+                                                      self.var_indice.next()),
+                                            size = expr.size)
+                setattr(identifier, self.__class__.is_var_ident, True)
                 self._vars[identifier] = expr
 
             # Recursion stop case
diff --git a/test/expression/expression_helper.py b/test/expression/expression_helper.py
index 3ff6f5da..514a9a51 100644
--- a/test/expression/expression_helper.py
+++ b/test/expression/expression_helper.py
@@ -30,11 +30,55 @@ class TestExpressionExpressionHelper(unittest.TestCase):
 
         # Test the result
         new_expr = vi.equation
+
         ## Force replace in the variable dependency order
         for var_id, var_value in reversed(vi.vars.items()):
             new_expr = new_expr.replace_expr({var_id: var_value})
         self.assertEqual(exprf, new_expr)
 
+        # Test prefix
+        vi = Variables_Identifier(exprf, var_prefix="prefix_v")
+
+        ## Use __str__
+        print vi
+
+        ## Test the result
+        new_expr = vi.equation
+        ### Force replace in the variable dependency order
+        for var_id, var_value in reversed(vi.vars.items()):
+            new_expr = new_expr.replace_expr({var_id: var_value})
+        self.assertEqual(exprf, new_expr)
+
+        # Test an identify on an expression already containing identifier
+        vi = Variables_Identifier(exprf)
+        vi2 = Variables_Identifier(vi.equation)
+
+        ## Test the result
+        new_expr = vi2.equation
+        ### Force replace in the variable dependency order
+        for var_id, var_value in reversed(vi2.vars.items()):
+            new_expr = new_expr.replace_expr({var_id: var_value})
+        self.assertEqual(vi.equation, new_expr)
+
+        ## Corner case: each sub var depends on itself
+        mem1 = m2_expr.ExprMem(ebx, size=32)
+        mem2 = m2_expr.ExprMem(mem1, size=32)
+        cst2 = m2_expr.ExprInt32(-1)
+        expr_mini = ((eax ^ mem2 ^ cst2) & (mem2 ^ (eax + mem2)))[31:32]
+
+        ## Build
+        vi = Variables_Identifier(expr_mini)
+        vi2 = Variables_Identifier(vi.equation)
+
+        ## Test the result
+        new_expr = vi2.equation
+        ### Force replace in the variable dependency order
+        for var_id, var_value in reversed(vi2.vars.items()):
+            new_expr = new_expr.replace_expr({var_id: var_value})
+        self.assertEqual(vi.equation, new_expr)
+
+
+
 if __name__ == '__main__':
     testcase = TestExpressionExpressionHelper
     testsuite = unittest.TestLoader().loadTestsFromTestCase(testcase)