Merge pull request #152 from commial/depgraph_z3

Depgraph z3

3 files changed, 120 insertions, 27 deletions

diff --git a/example/symbol_exec/depgraph.py b/example/symbol_exec/depgraph.py
index 802d4fca..5b6f373a 100644
--- a/example/symbol_exec/depgraph.py
+++ b/example/symbol_exec/depgraph.py
@@ -4,6 +4,7 @@ from pdb import pm
 from miasm2.analysis.machine import Machine
 from miasm2.analysis.binary import Container
 from miasm2.analysis.depgraph import DependencyGraph
+from miasm2.expression.expression import ExprMem, ExprId, ExprInt32
 
 parser = ArgumentParser("Dependency grapher")
 parser.add_argument("filename", help="Binary to analyse")
@@ -20,6 +21,9 @@ parser.add_argument("--unfollow-call", help="Stop on call statements",
 		    action="store_true")
 parser.add_argument("--do-not-simplify", help="Do not simplify expressions",
 		    action="store_true")
+parser.add_argument("--rename-args",
+                    help="Rename common arguments (@32[ESP_init] -> Arg1)",
+		    action="store_true")
 args = parser.parse_args()
 
 # Get architecture
@@ -41,6 +45,15 @@ for element in args.element:
 mdis = machine.dis_engine(cont.bin_stream, dont_dis_nulstart_bloc=True)
 ir_arch = machine.ira(mdis.symbol_pool)
 
+# Common argument forms
+init_ctx = {}
+if args.rename_args:
+    if arch == "x86_32":
+        # StdCall example
+        for i in xrange(4):
+            e_mem = ExprMem(ExprId("ESP_init") + ExprInt32(4 * (i + 1)), 32)
+            init_ctx[e_mem] = ExprId("arg%d" % i)
+
 # Disassemble the targeted function
 blocks = mdis.dis_multibloc(int(args.func_addr, 16))
 
@@ -71,7 +84,15 @@ for sol_nb, sol in enumerate(dg.get(current_block.label, elements, line_nb, set(
 	with open(fname, "w") as fdesc:
 		fdesc.write(sol.graph.dot())
 	result = ", ".join("%s: %s" % (k, v)
-			   for k, v in sol.emul().iteritems())
+			   for k, v in sol.emul(ctx=init_ctx).iteritems())
 	print "Solution %d: %s -> %s" % (sol_nb,
 					 result,
 					 fname)
+        if args.implicit:
+            sat = sol.is_satisfiable
+            constraints = ""
+            if sat:
+                constraints = {}
+                for element in sol.constraints:
+                    constraints[element] = hex(sol.constraints[element].as_long())
+            print "\tSatisfiability: %s %s" % (sat, constraints)
diff --git a/miasm2/analysis/depgraph.py b/miasm2/analysis/depgraph.py
index b92b3fd7..19f8e6e8 100644
--- a/miasm2/analysis/depgraph.py
+++ b/miasm2/analysis/depgraph.py
@@ -2,12 +2,18 @@
 import itertools
 from collections import namedtuple
 
+try:
+    import z3
+except ImportError:
+    pass
+
 import miasm2.expression.expression as m2_expr
 from miasm2.core.graph import DiGraph
 from miasm2.core.asmbloc import asm_label, expr_is_label
 from miasm2.expression.simplifications import expr_simp
 from miasm2.ir.symbexec import symbexec
 from miasm2.ir.ir import irbloc
+from miasm2.ir.translators import Translator
 
 class DependencyNode(object):
     """Node elements of a DependencyGraph
@@ -346,13 +352,18 @@ class DependencyResult(object):
     def input(self):
         return self._input_depnodes
 
-    def emul(self, step=False):
+    def emul(self, ctx=None, step=False):
         """Symbolic execution of relevant nodes according to the history
         Return the values of input nodes' elements
+        @ctx: (optional) Initial context as dictionnary
+        @step: (optional) Verbose execution
 
         /!\ The emulation is not safe if there is a loop in the relevant labels
         """
         # Init
+        ctx_init = self._ira.arch.regs.regs_init
+        if ctx is not None:
+            ctx_init.update(ctx)
         depnodes = self.relevant_nodes
         affects = []
 
@@ -366,7 +377,7 @@ class DependencyResult(object):
 
         # Eval the block
         temp_label = asm_label("Temp")
-        sb = symbexec(self._ira, self._ira.arch.regs.regs_init)
+        sb = symbexec(self._ira, ctx_init)
         sb.emulbloc(irbloc(temp_label, affects), step=step)
 
         # Return only inputs values (others could be wrongs)
@@ -374,6 +385,68 @@ class DependencyResult(object):
                 for depnode in self.input}
 
 
+class DependencyResultImplicit(DependencyResult):
+    """Stand for a result of a DependencyGraph with implicit option
+
+    Provide path constraints using the z3 solver"""
+
+    # Z3 Solver instance
+    _solver = None
+
+    def emul(self, ctx=None, step=False):
+        # Init
+        ctx_init = self._ira.arch.regs.regs_init
+        if ctx is not None:
+            ctx_init.update(ctx)
+        depnodes = self.relevant_nodes
+        solver = z3.Solver()
+        sb = symbexec(self._ira, ctx_init)
+        temp_label = asm_label("Temp")
+        history = self.relevant_labels[::-1]
+        history_size = len(history)
+
+        for hist_nb, label in enumerate(history):
+            # Build block with relevant lines only
+            affected_lines = set(depnode.line_nb for depnode in depnodes
+                                 if depnode.label == label)
+            irs = self._ira.blocs[label].irs
+            affects = []
+
+            for line_nb in sorted(affected_lines):
+                affects.append(irs[line_nb])
+
+            # Emul the block and get back destination
+            dst = sb.emulbloc(irbloc(temp_label, affects), step=step)
+
+            # Add constraint
+            if hist_nb + 1 < history_size:
+                next_label = history[hist_nb + 1]
+                expected = sb.eval_expr(m2_expr.ExprId(next_label, 32))
+                constraint = m2_expr.ExprAff(dst, expected)
+                solver.add(Translator.to_language("z3").from_expr(constraint))
+
+        # Save the solver
+        self._solver = solver
+
+        # Return only inputs values (others could be wrongs)
+        return {depnode.element: sb.symbols[depnode.element]
+                for depnode in self.input}
+
+    @property
+    def is_satisfiable(self):
+        """Return True iff the solution path admits at least one solution
+        PRE: 'emul'
+        """
+        return self._solver.check().r > 0
+
+    @property
+    def constraints(self):
+        """If satisfiable, return a valid solution as a Z3 Model instance"""
+        if not self.is_satisfiable:
+            raise ValueError("Unsatisfiable")
+        return self._solver.model()
+
+
 class FollowExpr(object):
     "Stand for an element (expression, depnode, ...) to follow or not"
 
@@ -458,11 +531,13 @@ class DependencyGraph(object):
     def _follow_label(exprs):
         """Do not follow labels"""
         follow = set()
+        unfollow = set()
         for expr in exprs:
             if expr_is_label(expr):
-                continue
-            follow.add(expr)
-        return follow, set()
+                unfollow.add(expr)
+            else:
+                follow.add(expr)
+        return follow, unfollow
 
     @staticmethod
     def _follow_mem_wrapper(exprs, mem_read):
@@ -546,7 +621,7 @@ class DependencyGraph(object):
 
             ## Build output
             output = FollowExpr.to_depnodes(read, depnode.label,
-                                                        depnode.line_nb - 1, modifier)
+                                            depnode.line_nb - 1, modifier)
 
         return output
 
@@ -588,12 +663,11 @@ class DependencyGraph(object):
 
     def _get_previousblocks(self, label):
         """Return an iterator on predecessors blocks of @label, with their
-        lengths and full block"""
+        lengths"""
         preds = self._ira.g.predecessors_iter(label)
         for pred_label in preds:
-            block = self._ira.blocs[pred_label]
-            length = len(block.irs)
-            yield (pred_label, length, block)
+            length = len(self._get_irs(pred_label))
+            yield (pred_label, length)
 
     def _processInterBloc(self, depnodes, heads):
         """Create a DependencyDict from @depnodes, and propagate DependencyDicts
@@ -631,12 +705,18 @@ class DependencyGraph(object):
             is_final = True
 
             # Propagate the DependencyDict to all parents
-            for label, irb_len, block in self._get_previousblocks(depdict.label):
+            for label, irb_len in self._get_previousblocks(depdict.label):
                 is_final = False
 
                 ## Duplicate the DependencyDict
                 new_depdict = depdict.extend(label)
 
+                if self._implicit:
+                    ### Implicit dependencies: IRDst will be link with heads
+                    implicit_depnode = DependencyNode(label, self._ira.IRDst,
+                                                      irb_len, modifier=False)
+                    new_depdict.pending.add(implicit_depnode)
+
                 ## Create links between DependencyDict
                 for depnode_head in depdict.pending:
                     ### Follow the head element in the parent
@@ -648,16 +728,7 @@ class DependencyGraph(object):
 
                     ### Handle implicit dependencies
                     if self._implicit:
-                        follow_exprs = self._follow_apply_cb(block.dst)
-                        fexpr_depnodes = FollowExpr.to_depnodes(follow_exprs,
-                                                                label,
-                                                                block.dst_linenb,
-                                                                False)
-                        extracted = FollowExpr.extract_depnodes(fexpr_depnodes)
-                        extfllw = FollowExpr.extract_depnodes(fexpr_depnodes,
-                                                              only_follow=True)
-                        new_depdict.cache[depnode_head].update(extracted)
-                        new_depdict.pending.update(extfllw)
+                        new_depdict.cache[depnode_head].add(implicit_depnode)
 
 
                 ## Manage the new element
@@ -689,6 +760,7 @@ class DependencyGraph(object):
 
         # Unify solutions
         unified = []
+        cls_res = DependencyResultImplicit if self._implicit else DependencyResult
         for final_depdict in depdicts:
             ## Keep only relevant nodes
             final_depdict.clean_modifiers_in_cache()
@@ -698,7 +770,7 @@ class DependencyGraph(object):
             if final_depdict not in unified:
                 unified.append(final_depdict)
                 ### Return solutions as DiGraph
-                yield DependencyResult(self._ira, final_depdict, input_depnodes)
+                yield cls_res(self._ira, final_depdict, input_depnodes)
 
     def get_fromDepNodes(self, depnodes, heads):
         """Alias for the get() method. Use the attributes of @depnodes as
diff --git a/test/test_all.py b/test/test_all.py
index 07e1c509..df07aac5 100644
--- a/test/test_all.py
+++ b/test/test_all.py
@@ -321,15 +321,15 @@ class ExampleSymbolExec(Example):
 
 
 testset += ExampleSymbolExec(["single_instr.py"])
-for options, nb_sol in [([], 8),
-                        (["-i"], 12)]:
+for options, nb_sol, tag in [([], 8, []),
+                             (["-i", "--rename-args"], 12, [TAGS["z3"]])]:
     testset += ExampleSymbolExec(["depgraph.py",
                                   Example.get_sample("simple_test.bin"),
                                   "-m", "x86_32", "0x0", "0x8b",
                                   "eax"] + options,
                                  products=["sol_%d.dot" % nb
-                                           for nb in xrange(nb_sol)])
-
+                                           for nb in xrange(nb_sol)],
+                                 tags=tag)
 
 ## Jitter
 class ExampleJitter(Example):