Merge pull request #310 from commial/feature-matchgraph

Feature MatchGraph

3 files changed, 329 insertions, 18 deletions

diff --git a/miasm2/core/asmbloc.py b/miasm2/core/asmbloc.py
index 9553d14d..aa26cfbf 100644
--- a/miasm2/core/asmbloc.py
+++ b/miasm2/core/asmbloc.py
@@ -10,7 +10,7 @@ import miasm2.expression.expression as m2_expr
 from miasm2.expression.simplifications import expr_simp
 from miasm2.expression.modint import moduint, modint
 from miasm2.core.utils import Disasm_Exception, pck
-from miasm2.core.graph import DiGraph, DiGraphSimplifier
+from miasm2.core.graph import DiGraph, DiGraphSimplifier, MatchGraphJoker
 from miasm2.core.interval import interval
 
 
@@ -980,27 +980,22 @@ class AsmCFG(DiGraph):
         if rebuild_needed:
             self.rebuild_edges()
 
-def _merge_blocks(dg, graph):
 
-    for block in graph.nodes():
+# Out of _merge_blocks to be computed only once
+_acceptable_block = lambda block: (not isinstance(block, asm_block_bad) and
+                                   len(block.lines) > 0)
+_parent = MatchGraphJoker(restrict_in=False, filt=_acceptable_block)
+_son = MatchGraphJoker(restrict_out=False, filt=_acceptable_block)
+_expgraph = _parent >> _son
 
-        # Acceptable block
-        if isinstance(block, asm_block_bad) or len(block.lines) == 0:
-            continue
 
-        # Only one son, not me
-        succs = graph.successors(block)
-        if len(succs) != 1:
-            continue
-        succ = succs[0]
-        if succ == block:
-            continue
+def _merge_blocks(dg, graph):
+    """Graph simplification merging asm_bloc with one and only one son with this
+    son if this son has one and only one parent"""
+    for match in _expgraph.match(graph):
 
-        # Son has only one parent, me
-        preds = graph.predecessors(succ)
-        if len(preds) != 1:
-            continue
-        assert preds[0] == block
+        # Get matching blocks
+        block, succ = match[_parent], match[_son]
 
         # Remove block last instruction if needed
         last_instr = block.lines[-1]
diff --git a/miasm2/core/graph.py b/miasm2/core/graph.py
index 66dffdba..310adc2e 100644
--- a/miasm2/core/graph.py
+++ b/miasm2/core/graph.py
@@ -572,3 +572,251 @@ class DiGraphSimplifier(object):
     def __call__(self, graph):
         """Wrapper on 'apply_simp'"""
         return self.apply_simp(graph)
+
+
+class MatchGraphJoker(object):
+    """MatchGraphJoker are joker nodes of MatchGraph, that is to say nodes which
+    stand for any node. Restrictions can be added to jokers.
+
+    If j1, j2 and j3 are MatchGraphJoker, one can quickly build a matcher for
+    the pattern:
+                                         |
+                                    +----v----+
+                                    |  (j1)   |
+                                    +----+----+
+                                         |
+                                    +----v----+
+                                    |  (j2)   |<---+
+                                    +----+--+-+    |
+                                         |  +------+
+                                    +----v----+
+                                    |  (j3)   |
+                                    +----+----+
+                                         |
+                                         v
+    Using:
+    >>> matcher = j1 >> j2 >> j3
+    >>> matcher += j2 >> j2
+    Or:
+    >>> matcher = j1 >> j2 >> j2 >> j3
+
+    """
+
+    def __init__(self, restrict_in=True, restrict_out=True, filt=None,
+                 name=None):
+        """Instanciate a MatchGraphJoker, with restrictions
+        @restrict_in: (optional) if set, the number of predecessors of the
+                      matched node must be the same than the joker node in the
+                      associated MatchGraph
+        @restrict_out: (optional) counterpart of @restrict_in for successors
+        @filt: (optional) function(node) -> boolean for filtering candidate node
+        @name: (optional) helper for displaying the current joker
+        """
+        if filt is None:
+            filt = lambda node: True
+        self.filt = filt
+        if name is None:
+            name = str(id(self))
+        self._name = name
+        self.restrict_in = restrict_in
+        self.restrict_out = restrict_out
+
+    def __rshift__(self, joker):
+        """Helper for describing a MatchGraph from @joker
+        J1 >> J2 stands for an edge going to J2 from J1
+        @joker: MatchGraphJoker instance
+        """
+        assert isinstance(joker, MatchGraphJoker)
+
+        graph = MatchGraph()
+        graph.add_node(self)
+        graph.add_node(joker)
+        graph.add_edge(self, joker)
+
+        # For future "A >> B" idiom construction
+        graph._last_node = joker
+
+        return graph
+
+    def __str__(self):
+        info = []
+        if not self.restrict_in:
+            info.append("In:*")
+        if not self.restrict_out:
+            info.append("Out:*")
+        return "Joker %s %s" % (self._name,
+                                "(%s)" % " ".join(info) if info else "")
+
+
+class MatchGraph(DiGraph):
+    """MatchGraph intends to be the counterpart of MatchExpr, but for DiGraph
+
+    This class provides API to match a given DiGraph pattern, with addidionnal
+    restrictions.
+    The implemented algorithm is a naive approach.
+
+    The recommended way to instanciate a MatchGraph is the use of
+    MatchGraphJoker.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(MatchGraph, self).__init__(*args, **kwargs)
+        # Construction helper
+        self._last_node = None
+
+    # Construction helpers
+    def __rshift__(self, joker):
+        """Construction helper, adding @joker to the current graph as a son of
+        _last_node
+        @joker: MatchGraphJoker instance"""
+        assert isinstance(joker, MatchGraphJoker)
+        assert isinstance(self._last_node, MatchGraphJoker)
+
+        self.add_node(joker)
+        self.add_edge(self._last_node, joker)
+        self._last_node = joker
+        return self
+
+    def __add__(self, graph):
+        """Construction helper, merging @graph with self
+        @graph: MatchGraph instance
+        """
+        assert isinstance(graph, MatchGraph)
+
+        # Reset helpers flag
+        self._last_node = None
+        graph._last_node = None
+
+        # Merge graph into self
+        for node in graph.nodes():
+            self.add_node(node)
+        for edge in graph.edges():
+            self.add_edge(*edge)
+
+        return self
+
+    # Graph matching
+    def _check_node(self, candidate, expected, graph, partial_sol=None):
+        """Check if @candidate can stand for @expected in @graph, given @partial_sol
+        @candidate: @graph's node
+        @expected: MatchGraphJoker instance
+        @graph: DiGraph instance
+        @partial_sol: (optional) dictionnary of MatchGraphJoker -> @graph's node
+        standing for a partial solution
+        """
+        # Avoid having 2 different joker for the same node
+        if partial_sol and candidate in partial_sol.values():
+            return False
+
+        # Check lambda filtering
+        if not expected.filt(candidate):
+            return False
+
+        # Check arity
+        ## If filter_in/out, then arity must be the same
+        ## Otherwise, arity of the candidate must be at least equal
+        if ((expected.restrict_in == True and
+             len(self.predecessors(expected)) != len(graph.predecessors(candidate))) or
+            (expected.restrict_in == False and
+             len(self.predecessors(expected)) > len(graph.predecessors(candidate)))):
+            return False
+        if ((expected.restrict_out == True and
+             len(self.successors(expected)) != len(graph.successors(candidate))) or
+            (expected.restrict_out == False and
+             len(self.successors(expected)) > len(graph.successors(candidate)))):
+            return False
+
+        # Check edges with partial solution if any
+        if not partial_sol:
+            return True
+        for pred in self.predecessors(expected):
+            if (pred in partial_sol and
+                partial_sol[pred] not in graph.predecessors(candidate)):
+                return False
+
+        for succ in self.successors(expected):
+            if (succ in partial_sol and
+                partial_sol[succ] not in graph.successors(candidate)):
+                return False
+
+        # All checks OK
+        return True
+
+    def _propagate_sol(self, node, partial_sol, graph, todo, propagator):
+        """
+        Try to extend the current @partial_sol by propagating the solution using
+        @propagator on @node.
+        New solutions are added to @todo
+        """
+        real_node = partial_sol[node]
+        for candidate in propagator(self, node):
+            ## Edge already in the partial solution, skip it
+            if candidate in partial_sol:
+                continue
+
+            ## Check candidate
+            for candidate_real in propagator(graph, real_node):
+                if self._check_node(candidate_real, candidate, graph,
+                                    partial_sol):
+                    temp_sol = partial_sol.copy()
+                    temp_sol[candidate] = candidate_real
+                    if temp_sol not in todo:
+                        todo.append(temp_sol)
+
+    @staticmethod
+    def _propagate_successors(graph, node):
+        """Propagate through @node successors in @graph"""
+        return graph.successors_iter(node)
+
+    @staticmethod
+    def _propagate_predecessors(graph, node):
+        """Propagate through @node predecessors in @graph"""
+        return graph.predecessors_iter(node)
+
+    def match(self, graph):
+        """Naive subgraph matching between graph and self.
+        Iterator on matching solution, as dictionnary MatchGraphJoker -> @graph
+        @graph: DiGraph instance
+        In order to obtained correct and complete results, @graph must be
+        connected.
+        """
+        # Partial solution: nodes corrects, edges between these nodes corrects
+        # A partial solution is a dictionnary MatchGraphJoker -> @graph's node
+        todo = list() # Dictionnaries containing partial solution
+        done = list() # Aleady computed partial solutions
+
+        # Elect first candidates
+        to_match = next(iter(self._nodes))
+        for node in graph.nodes():
+            if self._check_node(node, to_match, graph):
+                to_add = {to_match: node}
+                if to_add not in todo:
+                    todo.append(to_add)
+
+        while todo:
+            # When a partial_sol is computed, if more precise partial solutions
+            # are found, they will be added to 'todo'
+            # -> using last entry of todo first performs a "depth first"
+            # approach on solutions
+            # -> the algorithm may converge faster to a solution, a desired
+            # behavior while doing graph simplification (stopping after one sol)
+            partial_sol = todo.pop()
+
+            # Avoid infinite loop and recurrent work
+            if partial_sol in done:
+                continue
+            done.append(partial_sol)
+
+            # If all nodes are matching, this is a potential solution
+            if len(partial_sol) == len(self._nodes):
+                yield partial_sol
+                continue
+
+            # Find node to tests using edges
+            for node in partial_sol:
+                self._propagate_sol(node, partial_sol, graph, todo,
+                                    MatchGraph._propagate_successors)
+                self._propagate_sol(node, partial_sol, graph, todo,
+                                    MatchGraph._propagate_predecessors)
+
+        raise StopIteration
diff --git a/test/core/graph.py b/test/core/graph.py
index 33a2fc6f..e148d70f 100644
--- a/test/core/graph.py
+++ b/test/core/graph.py
@@ -217,3 +217,71 @@ for edge in graph3.edges():
     assert edge in graph4.edges()
 assert graph4.nodes() == graph.nodes().union(graph3.nodes())
 assert sorted(graph4.edges()) == sorted(graph.edges() + graph3.edges())
+
+# MatchGraph
+
+## Build a MatchGraph using MatchGraphJoker
+j1 = MatchGraphJoker(name="dad")
+j2 = MatchGraphJoker(name="son")
+### Check '>>' helper
+matcher = j1 >> j2 >> j1
+### Check __str__
+print matcher
+### Ensure form
+assert isinstance(matcher, MatchGraph)
+assert len(matcher.nodes()) == 2
+assert len(matcher.edges()) == 2
+
+## Match a simple graph
+graph = DiGraph()
+graph.add_edge(1, 2)
+graph.add_edge(2, 1)
+graph.add_edge(2, 3)
+sols = list(matcher.match(graph))
+assert len(sols) == 0
+
+## Modify restrictions
+j2 = MatchGraphJoker(name="son", restrict_out=False)
+matcher = j1 >> j2 >> j1
+sols = list(matcher.match(graph))
+assert len(sols) == 1
+assert sols[0] == {j1: 1,
+                   j2: 2}
+
+## Check solution combinaison (ie a -> b and b -> a)
+j1 = MatchGraphJoker(name="dad", restrict_out=False)
+matcher = j1 >> j2 >> j1
+sols = list(matcher.match(graph))
+assert len(sols) == 2
+assert len([sol for sol in sols if sol[j1] == 1]) == 1
+assert len([sol for sol in sols if sol[j1] == 2]) == 1
+
+## Check filter
+j2 = MatchGraphJoker(name="son", restrict_out=False, filt=lambda node: node < 2)
+matcher = j1 >> j2 >> j1
+sols = list(matcher.match(graph))
+assert len(sols) == 1
+assert sols[0] == {j1: 2,
+                   j2: 1}
+
+## Check building with 'add' helper
+j1 = MatchGraphJoker(name="dad")
+j2 = MatchGraphJoker(name="son")
+j3 = MatchGraphJoker(name="sonson", restrict_in=False)
+matcher = j1 >> j2
+matcher += j2 >> j3
+assert isinstance(matcher, MatchGraph)
+assert len(matcher.nodes()) == 3
+assert len(matcher.edges()) == 2
+
+## Check restrict_in
+graph = DiGraph()
+graph.add_edge(1, 2)
+graph.add_edge(2, 3)
+graph.add_edge(4, 3)
+sols = list(matcher.match(graph))
+assert len(sols) == 1
+assert sols[0] == {j1: 1,
+                   j2: 2,
+                   j3: 3}
+