Merge pull request #123 from serpilliere/depgraph_filter

Depgraph filter

3 files changed, 106 insertions, 57 deletions

diff --git a/example/ida/depgraph.py b/example/ida/depgraph.py
index edf4ce0b..ae00c357 100644
--- a/example/ida/depgraph.py
+++ b/example/ida/depgraph.py
@@ -9,7 +9,7 @@ from miasm2.expression import expression as m2_expr
 
 from miasm2.expression.simplifications import expr_simp
 from miasm2.analysis.machine import Machine
-from miasm2.analysis.depgraph import DependencyGraph, DependencyGraph_NoMemory
+from miasm2.analysis.depgraph import DependencyGraph
 
 from utils import guess_machine
 
@@ -49,8 +49,8 @@ Track the element:
 <Line number:{iLineNb}>
 
 Method to use:
-<Best effort:{rBest}>
-<No memory (sound & complete):{rNoMem}>{cMethod}>
+<Follow Memory:{rNoMem}>
+<Follow Call:{rNoCall}>{cMethod}>
 
 <Highlight color:{cColor}>
 """, {
@@ -60,7 +60,7 @@ Method to use:
                     selval=reg_default),
             'cMode': Form.RadGroupControl(("rBeforeLine", "rAfterLine",
                                            "rEndBlock")),
-            'cMethod': Form.RadGroupControl(("rBest", "rNoMem")),
+            'cMethod': Form.ChkGroupControl(("rNoMem", "rNoCall")),
             'iLineNb': Form.NumericInput(tp=Form.FT_RAWHEX,
                                          value=line_nb),
             'cbBBL': Form.DropdownListControl(
@@ -97,14 +97,11 @@ Method to use:
         return set([ir_arch.arch.regs.all_regs_ids_byname[value]])
 
     @property
-    def method(self):
+    def depgraph(self):
         value = self.cMethod.value
-        if value == 0:
-            return DependencyGraph
-        elif value == 1:
-            return DependencyGraph_NoMemory
-        else:
-            raise ValueError("Unknown method")
+        return DependencyGraph(self.ira,
+                               follow_mem=value & 1,
+                               follow_call=value & 2)
 
     @property
     def color(self):
@@ -148,7 +145,7 @@ settings = depGraphSettingsForm(ir_arch)
 settings.Execute()
 
 # Get dependency graphs
-dg = (settings.method)(ir_arch)
+dg = settings.depgraph
 graphs = dg.get(settings.label, settings.elements, settings.line_nb,
                 set([ir_arch.symbol_pool.getby_offset(func.startEA)]))
 
diff --git a/miasm2/analysis/depgraph.py b/miasm2/analysis/depgraph.py
index ead5c8a9..1dca9fb3 100644
--- a/miasm2/analysis/depgraph.py
+++ b/miasm2/analysis/depgraph.py
@@ -1,5 +1,6 @@
 """Provide dependency graph"""
 import itertools
+from collections import namedtuple
 
 import miasm2.expression.expression as m2_expr
 from miasm2.core.graph import DiGraph
@@ -8,7 +9,6 @@ from miasm2.expression.simplifications import expr_simp
 from miasm2.ir.symbexec import symbexec
 from miasm2.ir.ir import irbloc
 
-
 class DependencyNode(object):
     """Node elements of a DependencyGraph
 
@@ -374,6 +374,8 @@ class DependencyResult(object):
                 for depnode in self.input}
 
 
+FollowExpr = namedtuple("FollowExpr", ["follow", "element"])
+
 class DependencyGraph(object):
     """Implementation of a dependency graph
 
@@ -383,15 +385,86 @@ class DependencyGraph(object):
     *explicitely* involved in the equation of given element.
     """
 
-    def __init__(self, ira):
+    def __init__(self, ira, 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
+
+        Following arguments define filters used to generate dependencies
+        @apply_simp: (optional) Apply expr_simp
+        @follow_mem: (optional) Track memory syntaxically
+        @follow_call: (optional) Track throught "call"
         """
         # Init
         self._ira = ira
 
         # The IRA graph must be computed
-        self._ira.gen_graph()
+        assert(hasattr(self._ira, 'g'))
+
+        # Create callback filters. The order is relevant.
+        self._cb_follow = []
+        if apply_simp:
+            self._cb_follow.append(self._follow_simp_expr)
+        if follow_mem:
+            self._cb_follow.append(self._follow_mem)
+        else:
+            self._cb_follow.append(self._follow_nomem)
+        if not follow_call:
+            self._cb_follow.append(self._follow_nocall)
+
+    @staticmethod
+    def _follow_simp_expr(exprs):
+        """Simplify expression so avoid tracking useless elements,
+        as: XOR EAX, EAX
+        """
+        follow = set()
+        for expr in exprs:
+            follow.add(expr_simp(expr))
+        return follow, set()
+
+    @staticmethod
+    def _follow_mem_wrapper(exprs, mem_read):
+        follow = set()
+        for expr in exprs:
+            follow.update(expr.get_r(mem_read=mem_read, cst_read=True))
+        return follow, set()
+
+    @staticmethod
+    def _follow_mem(exprs):
+        """Follow expression from memory pointer"""
+        return DependencyGraph._follow_mem_wrapper(exprs, True)
+
+    @staticmethod
+    def _follow_nomem(exprs):
+        """Don't follow expression from memory pointer"""
+        return DependencyGraph._follow_mem_wrapper(exprs, False)
+
+    @staticmethod
+    def _follow_nocall(exprs):
+        """Don't follow expression from sub_call"""
+        follow = set()
+        nofollow = set()
+        for expr in exprs:
+            if isinstance(expr, m2_expr.ExprOp) and expr.op.startswith('call'):
+                nofollow.add(expr)
+            else:
+                follow.add(expr)
+        return follow, nofollow
+
+    def _follow_apply_cb(self, expr):
+        follow = set([expr])
+        nofollow = set()
+
+        for cb in self._cb_follow:
+            follow, nofollow_tmp = cb(follow)
+            nofollow.update(nofollow_tmp)
+
+        out = set(FollowExpr(True, expr) for expr in follow)
+        out.update(set(FollowExpr(False, expr) for expr in nofollow))
+
+        return out
 
     def _get_irs(self, label):
         "Return the irs associated to @label"
@@ -403,7 +476,8 @@ class DependencyGraph(object):
         return self._get_irs(depnode.label)[depnode.line_nb - 1]
 
     def _resolve_depNode(self, depnode):
-        """Compute and return the dependencies involved by @depnode"""
+        """Compute and return the dependencies involved by @depnode
+        Return a set of FollowExpr"""
 
         if isinstance(depnode.element, m2_expr.ExprInt):
             # A constant does not have any dependency
@@ -417,27 +491,27 @@ class DependencyGraph(object):
             # Intra-block resolving
             ## Get dependencies
             read = set()
+
             modifier = False
 
             for affect in self._get_affblock(depnode):
                 if affect.dst == depnode.element:
-                    ### Avoid tracking useless elements, as XOR EAX, EAX
-                    src = expr_simp(affect.src)
-
-                    read.update(src.get_r(mem_read=True, cst_read=True))
+                    elements = self._follow_apply_cb(affect.src)
+                    read.update(elements)
                     modifier = True
 
             ## If it's not a modifier affblock, reinject current element
             if not modifier:
-                read = set([depnode.element])
+                read = set([FollowExpr(True, depnode.element)])
 
             ## Build output
             dependencies = set()
-            for element in read:
-                dependencies.add(DependencyNode(depnode.label,
-                                                element,
-                                                depnode.line_nb - 1,
-                                                modifier=modifier))
+            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
 
         return output
@@ -465,10 +539,13 @@ class DependencyGraph(object):
 
             # Find dependency of the current depnode
             sub_depnodes = self._resolve_depNode(depnode)
-            depdict.cache[depnode] = sub_depnodes
+            depdict.cache[depnode] = set(follow_expr.element
+                                         for follow_expr in sub_depnodes)
 
             # Add to the worklist its dependencies
-            todo.update(sub_depnodes)
+            todo.update(set(follow_expr.element
+                            for follow_expr in sub_depnodes
+                            if follow_expr.follow))
 
         # Pending states will be override in cache
         for depnode in depdict.pending:
@@ -596,24 +673,3 @@ class DependencyGraph(object):
         """
         return self.get(label, elements, len(self._get_irs(label)), heads)
 
-
-class DependencyGraph_NoMemory(DependencyGraph):
-    """Dependency graph without memory tracking.
-
-    That way, the output has following properties:
-    - Only explicit dependencies are followed
-    - Soundness: all results are corrects
-    - Completeness: all possible solutions are founds
-    """
-
-    def _resolve_depNode(self, depnode):
-        """Compute and return the dependencies involved by @depnode"""
-        # Get inital elements
-        result = super(DependencyGraph_NoMemory, self)._resolve_depNode(depnode)
-
-        # If @depnode depends on a memory element, give up
-        for node in result:
-            if isinstance(node.element, m2_expr.ExprMem):
-                return set()
-
-        return result
diff --git a/test/analysis/depgraph.py b/test/analysis/depgraph.py
index 5484ec02..a661d785 100644
--- a/test/analysis/depgraph.py
+++ b/test/analysis/depgraph.py
@@ -3,7 +3,7 @@ from miasm2.core.asmbloc import asm_label
 from miasm2.ir.analysis import ira
 from miasm2.ir.ir import ir, irbloc
 from miasm2.core.graph import DiGraph
-from miasm2.analysis.depgraph import DependencyNode, DependencyGraph, DependencyDict, DependencyGraph_NoMemory
+from miasm2.analysis.depgraph import DependencyNode, DependencyGraph, DependencyDict
 from pdb import pm
 
 a = ExprId("a")
@@ -58,9 +58,6 @@ class IRATest(ir, ira):
         ir.__init__(self, arch, 32, symbol_pool)
         self.IRDst = pc
 
-    def gen_graph(self):
-        return
-
 class GraphTest(DiGraph):
     def __init__(self, ira):
         self.ira = ira
@@ -73,9 +70,6 @@ class GraphTest(DiGraph):
             return False
         return True
 
-    def gen_graph(self):
-        return
-
     def node2str(self, node):
         if not node in self.ira.blocs:
             return str(node)
@@ -603,7 +597,9 @@ for i, test in enumerate([(g1_ira, g1_input, [g1_output1]),
     open("graph_%02d.dot" % (i+1), "w").write(g_ira.g.dot())
     # Test classes
     for g_dep in [DependencyGraph(g_ira),
-                  DependencyGraph_NoMemory(g_ira)]:
+                  DependencyGraph(g_ira, apply_simp=False),
+                  DependencyGraph(g_ira, follow_mem=False),
+                  DependencyGraph(g_ira, follow_mem=False, follow_call=False)]:
         print " - Class %s" % g_dep.__class__.__name__
 
         ## Test public APIs