Merge pull request #148 from commial/depgraph

Depgraph Implicit

7 files changed, 183 insertions, 40 deletions

diff --git a/example/ida/depgraph.py b/example/ida/depgraph.py
index ae00c357..fcd87cb9 100644
--- a/example/ida/depgraph.py
+++ b/example/ida/depgraph.py
@@ -50,7 +50,8 @@ Track the element:
 
 Method to use:
 <Follow Memory:{rNoMem}>
-<Follow Call:{rNoCall}>{cMethod}>
+<Follow Call:{rNoCall}>
+<Implicit dependencies:{rImplicit}>{cMethod}>
 
 <Highlight color:{cColor}>
 """, {
@@ -60,7 +61,7 @@ Method to use:
                     selval=reg_default),
             'cMode': Form.RadGroupControl(("rBeforeLine", "rAfterLine",
                                            "rEndBlock")),
-            'cMethod': Form.ChkGroupControl(("rNoMem", "rNoCall")),
+            'cMethod': Form.ChkGroupControl(("rNoMem", "rNoCall", "rImplicit")),
             'iLineNb': Form.NumericInput(tp=Form.FT_RAWHEX,
                                          value=line_nb),
             'cbBBL': Form.DropdownListControl(
@@ -100,6 +101,7 @@ Method to use:
     def depgraph(self):
         value = self.cMethod.value
         return DependencyGraph(self.ira,
+                               implicit=value & 4,
                                follow_mem=value & 1,
                                follow_call=value & 2)
 
diff --git a/example/symbol_exec/depgraph.py b/example/symbol_exec/depgraph.py
new file mode 100644
index 00000000..802d4fca
--- /dev/null
+++ b/example/symbol_exec/depgraph.py
@@ -0,0 +1,77 @@
+from argparse import ArgumentParser
+from pdb import pm
+
+from miasm2.analysis.machine import Machine
+from miasm2.analysis.binary import Container
+from miasm2.analysis.depgraph import DependencyGraph
+
+parser = ArgumentParser("Dependency grapher")
+parser.add_argument("filename", help="Binary to analyse")
+parser.add_argument("func_addr", help="Function address")
+parser.add_argument("target_addr", help="Address to start")
+parser.add_argument("element", nargs="+", help="Elements to track")
+parser.add_argument("-m", "--architecture",
+		    help="Architecture (%s)" % Machine.available_machine())
+parser.add_argument("-i", "--implicit", help="Use implicit tracking",
+		    action="store_true")
+parser.add_argument("--unfollow-mem", help="Stop on memory statements",
+		    action="store_true")
+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")
+args = parser.parse_args()
+
+# Get architecture
+with open(args.filename) as fstream:
+    cont = Container.from_stream(fstream)
+
+arch = args.architecture if args.architecture else cont.arch
+machine = Machine(arch)
+
+# Check elements
+elements = set()
+regs = machine.mn.regs.all_regs_ids_byname
+for element in args.element:
+    try:
+	elements.add(regs[element.upper()])
+    except KeyError:
+	raise ValueError("Unknown element '%s'" % element)
+
+mdis = machine.dis_engine(cont.bin_stream, dont_dis_nulstart_bloc=True)
+ir_arch = machine.ira(mdis.symbol_pool)
+
+# Disassemble the targeted function
+blocks = mdis.dis_multibloc(int(args.func_addr, 16))
+
+# Generate IR
+for block in blocks:
+    ir_arch.add_bloc(block)
+
+# Build the IRA Graph
+ir_arch.gen_graph()
+
+# Get the instance
+dg = DependencyGraph(ir_arch, implicit=args.implicit,
+		     apply_simp=not(args.do_not_simplify),
+		     follow_mem=not(args.unfollow_mem),
+		     follow_call=not(args.unfollow_call))
+
+# Build information
+target_addr = int(args.target_addr, 16)
+current_block = list(ir_arch.getby_offset(target_addr))[0]
+line_nb = 0
+for line_nb, line in enumerate(current_block.lines):
+    if line.offset == target_addr:
+	break
+
+# Enumerate solutions
+for sol_nb, sol in enumerate(dg.get(current_block.label, elements, line_nb, set())):
+	fname = "sol_%d.dot" % sol_nb
+	with open(fname, "w") as fdesc:
+		fdesc.write(sol.graph.dot())
+	result = ", ".join("%s: %s" % (k, v)
+			   for k, v in sol.emul().iteritems())
+	print "Solution %d: %s -> %s" % (sol_nb,
+					 result,
+					 fname)
diff --git a/miasm2/analysis/depgraph.py b/miasm2/analysis/depgraph.py
index 1dca9fb3..b92b3fd7 100644
--- a/miasm2/analysis/depgraph.py
+++ b/miasm2/analysis/depgraph.py
@@ -4,7 +4,7 @@ from collections import namedtuple
 
 import miasm2.expression.expression as m2_expr
 from miasm2.core.graph import DiGraph
-from miasm2.core.asmbloc import asm_label
+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
@@ -374,7 +374,34 @@ class DependencyResult(object):
                 for depnode in self.input}
 
 
-FollowExpr = namedtuple("FollowExpr", ["follow", "element"])
+class FollowExpr(object):
+    "Stand for an element (expression, depnode, ...) to follow or not"
+
+    def __init__(self, follow, element):
+        self.follow = follow
+        self.element = element
+
+    @staticmethod
+    def to_depnodes(follow_exprs, label, line, modifier):
+        """Build a set of FollowExpr(DependencyNode) from the @follow_exprs set
+        of FollowExpr"""
+        dependencies = set()
+        for follow_expr in follow_exprs:
+            dependencies.add(FollowExpr(follow_expr.follow,
+                                        DependencyNode(label,
+                                                       follow_expr.element,
+                                                       line,
+                                                       modifier=modifier)))
+        return dependencies
+
+    @staticmethod
+    def extract_depnodes(follow_exprs, only_follow=False):
+        """Extract depnodes from a set of FollowExpr(Depnodes)
+        @only_follow: (optional) extract only elements to follow"""
+        return set(follow_expr.element
+                   for follow_expr in follow_exprs
+                   if not(only_follow) or follow_expr.follow)
+
 
 class DependencyGraph(object):
     """Implementation of a dependency graph
@@ -382,15 +409,16 @@ class DependencyGraph(object):
     A dependency graph contains DependencyNode as nodes. The oriented edges
     stand for a dependency.
     The dependency graph is made of the lines of a group of IRblock
-    *explicitely* involved in the equation of given element.
+    *explicitely* or *implicitely* involved in the equation of given element.
     """
 
-    def __init__(self, ira, apply_simp=True, follow_mem=True,
+    def __init__(self, ira, implicit=False, apply_simp=True, follow_mem=True,
                  follow_call=True):
         """Create a DependencyGraph linked to @ira
         The IRA graph must have been computed
 
         @ira: IRAnalysis instance
+        @implicit: (optional) Imply implicit dependencies
 
         Following arguments define filters used to generate dependencies
         @apply_simp: (optional) Apply expr_simp
@@ -399,6 +427,7 @@ class DependencyGraph(object):
         """
         # Init
         self._ira = ira
+        self._implicit = implicit
 
         # The IRA graph must be computed
         assert(hasattr(self._ira, 'g'))
@@ -413,6 +442,7 @@ class DependencyGraph(object):
             self._cb_follow.append(self._follow_nomem)
         if not follow_call:
             self._cb_follow.append(self._follow_nocall)
+        self._cb_follow.append(self._follow_label)
 
     @staticmethod
     def _follow_simp_expr(exprs):
@@ -425,6 +455,16 @@ class DependencyGraph(object):
         return follow, set()
 
     @staticmethod
+    def _follow_label(exprs):
+        """Do not follow labels"""
+        follow = set()
+        for expr in exprs:
+            if expr_is_label(expr):
+                continue
+            follow.add(expr)
+        return follow, set()
+
+    @staticmethod
     def _follow_mem_wrapper(exprs, mem_read):
         follow = set()
         for expr in exprs:
@@ -505,14 +545,8 @@ class DependencyGraph(object):
                 read = set([FollowExpr(True, depnode.element)])
 
             ## Build output
-            dependencies = set()
-            for follow_expr in read:
-                dependencies.add(FollowExpr(follow_expr.follow,
-                                            DependencyNode(depnode.label,
-                                                           follow_expr.element,
-                                                           depnode.line_nb - 1,
-                                                           modifier=modifier)))
-            output = dependencies
+            output = FollowExpr.to_depnodes(read, depnode.label,
+                                                        depnode.line_nb - 1, modifier)
 
         return output
 
@@ -539,13 +573,11 @@ class DependencyGraph(object):
 
             # Find dependency of the current depnode
             sub_depnodes = self._resolve_depNode(depnode)
-            depdict.cache[depnode] = set(follow_expr.element
-                                         for follow_expr in sub_depnodes)
+            depdict.cache[depnode] = FollowExpr.extract_depnodes(sub_depnodes)
 
             # Add to the worklist its dependencies
-            todo.update(set(follow_expr.element
-                            for follow_expr in sub_depnodes
-                            if follow_expr.follow))
+            todo.update(FollowExpr.extract_depnodes(sub_depnodes,
+                                                    only_follow=True))
 
         # Pending states will be override in cache
         for depnode in depdict.pending:
@@ -556,11 +588,12 @@ class DependencyGraph(object):
 
     def _get_previousblocks(self, label):
         """Return an iterator on predecessors blocks of @label, with their
-        lengths"""
+        lengths and full block"""
         preds = self._ira.g.predecessors_iter(label)
         for pred_label in preds:
-            length = len(self._get_irs(pred_label))
-            yield (pred_label, length)
+            block = self._ira.blocs[pred_label]
+            length = len(block.irs)
+            yield (pred_label, length, block)
 
     def _processInterBloc(self, depnodes, heads):
         """Create a DependencyDict from @depnodes, and propagate DependencyDicts
@@ -598,7 +631,7 @@ class DependencyGraph(object):
             is_final = True
 
             # Propagate the DependencyDict to all parents
-            for label, irb_len in self._get_previousblocks(depdict.label):
+            for label, irb_len, block in self._get_previousblocks(depdict.label):
                 is_final = False
 
                 ## Duplicate the DependencyDict
@@ -613,6 +646,20 @@ class DependencyGraph(object):
                     new_depdict.cache[depnode_head] = set([new_depnode])
                     new_depdict.pending.add(new_depnode)
 
+                    ### 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)
+
+
                 ## Manage the new element
                 todo.append(new_depdict)
 
diff --git a/miasm2/arch/arm/ira.py b/miasm2/arch/arm/ira.py
index 74548f86..b918a2e6 100644
--- a/miasm2/arch/arm/ira.py
+++ b/miasm2/arch/arm/ira.py
@@ -69,7 +69,7 @@ class ir_a_arml(ir_a_arml_base):
             nbloc = irbloc(new_lbl, irs)
             nbloc.lines = [l]*len(irs)
             self.blocs[new_lbl] = nbloc
-            irb.set_dst(ExprId(new_lbl, size=self.pc.size))
+            irb.dst = ExprId(new_lbl, size=self.pc.size)
 
         """
         if not bloc.lines:
diff --git a/miasm2/ir/ir.py b/miasm2/ir/ir.py
index 252f0ab3..5d77c5a1 100644
--- a/miasm2/ir/ir.py
+++ b/miasm2/ir/ir.py
@@ -36,34 +36,42 @@ class irbloc(object):
         self.lines = lines
         self.except_automod = True
         self._dst = None
+        self._dst_linenb = None
 
 
-    def get_dst(self):
+    def _get_dst(self):
+        """Find the IRDst affectation and update dst, dst_linenb accordingly"""
         if self._dst is not None:
             return self._dst
         dst = None
-        for ir in self.irs:
+        for linenb, ir in enumerate(self.irs):
             for i in ir:
                 if isinstance(i.dst, m2_expr.ExprId) and i.dst.name == "IRDst":
                     if dst is not None:
                         raise ValueError('Multiple destinations!')
                     dst = i.src
+                    dst_linenb = linenb
         self._dst = dst
+        self._dst_linenb = linenb
         return dst
 
-    def set_dst(self, value):
+    def _set_dst(self, value):
         """Find and replace the IRDst affectation's source by @value"""
-        dst = None
-        for ir in self.irs:
-            for i, expr in enumerate(ir):
-                if isinstance(expr.dst, m2_expr.ExprId) and expr.dst.name == "IRDst":
-                    if dst is not None:
-                        raise ValueError('Multiple destinations!')
-                    dst = value
-                    ir[i] = m2_expr.ExprAff(expr.dst, value)
+        if self._dst_linenb is None:
+            self._get_dst()
+
+        ir = self.irs[self._dst_linenb]
+        for i, expr in enumerate(ir):
+            if isinstance(expr.dst, m2_expr.ExprId) and expr.dst.name == "IRDst":
+                ir[i] = m2_expr.ExprAff(expr.dst, value)
         self._dst = value
 
-    dst = property(get_dst, set_dst)
+    dst = property(_get_dst, _set_dst)
+
+    @property
+    def dst_linenb(self):
+        """Line number of the IRDst setting statement in the current irs"""
+        return self._dst_linenb
 
     def get_rw(self):
         self.r = []
diff --git a/test/analysis/depgraph.py b/test/analysis/depgraph.py
index a661d785..c5e41f64 100644
--- a/test/analysis/depgraph.py
+++ b/test/analysis/depgraph.py
@@ -615,15 +615,16 @@ for i, test in enumerate([(g1_ira, g1_input, [g1_output1]),
             g_list = list(g_list)
             ### Dump outputs graphs for debug means
             for j, result in enumerate(g_list):
-                open("graph_test_%02d_%02d.dot" % (i+1, j), "w").write(result.graph.dot())
+                open("graph_test_%02d_%02d.dot" % (i + 1, j),
+                     "w").write(result.graph.dot())
 
             ### The number of results should be the same
             print " - - - number of results"
             assert(len(g_list) == len(g_test_list))
 
             ### Match the right result (unordered)
-            for i, result in enumerate(g_list):
-                print " - - - result %d" % i
+            for j, result in enumerate(g_list):
+                print " - - - result %d" % j
                 found = False
                 for expected in g_test_list:
                     if expected["graph"].__eq__(result.graph):
diff --git a/test/test_all.py b/test/test_all.py
index aefd0196..07e1c509 100644
--- a/test/test_all.py
+++ b/test/test_all.py
@@ -128,7 +128,6 @@ for script in ["win_api_x86_32.py",
 testset += RegressionTest(["depgraph.py"], base_dir="analysis",
                           products=["graph_test_01_00.dot",
                                     "graph_test_02_00.dot",
-                                    "graph_test_02_01.dot",
                                     "graph_test_03_00.dot",
                                     "graph_test_03_01.dot",
                                     "graph_test_04_00.dot",
@@ -322,6 +321,15 @@ class ExampleSymbolExec(Example):
 
 
 testset += ExampleSymbolExec(["single_instr.py"])
+for options, nb_sol in [([], 8),
+                        (["-i"], 12)]:
+    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)])
+
 
 ## Jitter
 class ExampleJitter(Example):