about summary refs log tree commit diff stats
path: root/miasm2
diff options
context:
space:
mode:
Diffstat (limited to 'miasm2')
-rw-r--r--miasm2/analysis/depgraph.py87
-rw-r--r--miasm2/ir/ir.py32
2 files changed, 72 insertions, 47 deletions
diff --git a/miasm2/analysis/depgraph.py b/miasm2/analysis/depgraph.py
index 8d6da3b2..838183bf 100644
--- a/miasm2/analysis/depgraph.py
+++ b/miasm2/analysis/depgraph.py
@@ -16,6 +16,7 @@ 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
@@ -27,6 +28,7 @@ class DependencyNode(object):
 
     __slots__ = ["_label", "_element", "_line_nb", "_modifier",
                  "_step", "_nostep_repr", "_hash"]
+
     def __init__(self, label, element, line_nb, step, modifier=False):
         """Create a dependency node with:
         @label: asm_label instance
@@ -299,7 +301,6 @@ class DependencyDict(object):
 
         self._cache = CacheWrapper(self._get_modifiers_in_cache(node_heads))
 
-
     def _build_depgraph(self, depnode):
         """Recursively build the final list of DiGraph, and clean up unmodifier
         nodes
@@ -495,7 +496,7 @@ class DependencyResult(object):
     def unresolved(self):
         """Set of nodes whose dependencies weren't found"""
         return set(node.nostep_repr for node in self._depdict.pending
-                    if node.element != self._ira.IRDst)
+                   if node.element != self._ira.IRDst)
 
     @property
     def relevant_nodes(self):
@@ -708,13 +709,10 @@ class DependencyGraph(object):
         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)
-        self._cb_follow.append(self._follow_label)
+        self._cb_follow.append(lambda exprs: self._follow_exprs(exprs,
+                                                                follow_mem,
+                                                                follow_call))
+        self._cb_follow.append(self._follow_nolabel)
 
     @property
     def step_counter(self):
@@ -742,44 +740,59 @@ class DependencyGraph(object):
         return follow, set()
 
     @staticmethod
-    def _follow_label(exprs):
-        """Do not follow labels"""
-        follow = set()
-        for expr in exprs:
-            if not expr_is_label(expr):
-                follow.add(expr)
+    def get_expr(expr, follow, nofollow):
+        """Update @follow/@nofollow according to insteresting nodes
+        Returns same expression (non modifier visitor).
 
-        return follow, set()
+        @expr: expression to handle
+        @follow: set of nodes to follow
+        @nofollow: set of nodes not to follow
+        """
+        if isinstance(expr, m2_expr.ExprId):
+            follow.add(expr)
+        elif isinstance(expr, m2_expr.ExprInt):
+            nofollow.add(expr)
+        return expr
 
     @staticmethod
-    def _follow_mem_wrapper(exprs, mem_read):
-        """Wrapper to follow or not expression from memory pointer"""
-        follow = set()
-        for expr in exprs:
-            follow.update(expr.get_r(mem_read=mem_read, cst_read=True))
-        return follow, set()
+    def follow_expr(expr, follow, nofollow, follow_mem=False, follow_call=False):
+        """Returns True if we must visit sub expressions.
+        @expr: expression to browse
+        @follow: set of nodes to follow
+        @nofollow: set of nodes not to follow
+        @follow_mem: force the visit of memory sub expressions
+        @follow_call: force the visit of call sub expressions
+        """
+        if not follow_mem and isinstance(expr, m2_expr.ExprMem):
+            nofollow.add(expr)
+            return False
+        if not follow_call and expr.is_function_call():
+            nofollow.add(expr)
+            return False
+        return True
 
-    @staticmethod
-    def _follow_mem(exprs):
-        """Follow expression from memory pointer"""
-        return DependencyGraph._follow_mem_wrapper(exprs, True)
+    @classmethod
+    def _follow_exprs(cls, exprs, follow_mem=False, follow_call=False):
+        """Extracts subnodes from exprs and returns followed/non followed
+        expressions according to @follow_mem/@follow_call
 
-    @staticmethod
-    def _follow_nomem(exprs):
-        """Don't follow expression from memory pointer"""
-        return DependencyGraph._follow_mem_wrapper(exprs, False)
+        """
+        follow, nofollow = set(), set()
+        for expr in exprs:
+            expr.visit(lambda x: cls.get_expr(x, follow, nofollow),
+                       lambda x: cls.follow_expr(x, follow, nofollow,
+                                                 follow_mem, follow_call))
+        return follow, nofollow
 
     @staticmethod
-    def _follow_nocall(exprs):
-        """Don't follow expression from sub_call"""
+    def _follow_nolabel(exprs):
+        """Do not follow labels"""
         follow = set()
-        nofollow = set()
         for expr in exprs:
-            if expr.is_function_call():
-                nofollow.add(expr)
-            else:
+            if not expr_is_label(expr):
                 follow.add(expr)
-        return follow, nofollow
+
+        return follow, set()
 
     def _follow_apply_cb(self, expr):
         """Apply callback functions to @expr
diff --git a/miasm2/ir/ir.py b/miasm2/ir/ir.py
index 32c97661..e051dc8c 100644
--- a/miasm2/ir/ir.py
+++ b/miasm2/ir/ir.py
@@ -135,17 +135,27 @@ class ir(object):
         ir_bloc_cur, ir_blocs_extra = self.get_ir(l)
         return ir_bloc_cur, ir_blocs_extra
 
-    def get_bloc(self, ad):
-        if isinstance(ad, m2_expr.ExprId) and isinstance(ad.name,
-                                                         asmbloc.asm_label):
+    def get_label(self, ad):
+        """Transforms an ExprId/ExprInt/label/int into a label
+        @ad: an ExprId/ExprInt/label/int"""
+
+        if (isinstance(ad, m2_expr.ExprId) and
+            isinstance(ad.name, asmbloc.asm_label)):
             ad = ad.name
         if isinstance(ad, m2_expr.ExprInt):
             ad = int(ad.arg)
         if type(ad) in [int, long]:
-            ad = self.symbol_pool.getby_offset(ad)
+            ad = self.symbol_pool.getby_offset_create(ad)
         elif isinstance(ad, asmbloc.asm_label):
-            ad = self.symbol_pool.getby_name(ad.name)
-        return self.blocs.get(ad, None)
+            ad = self.symbol_pool.getby_name_create(ad.name)
+        return ad
+
+    def get_bloc(self, ad):
+        """Returns the irbloc associated to an ExprId/ExprInt/label/int
+        @ad: an ExprId/ExprInt/label/int"""
+
+        label = self.get_label(ad)
+        return self.blocs.get(label, None)
 
     def add_instr(self, l, ad=0, gen_pc_updt = False):
         b = asmbloc.asm_bloc(l)
@@ -227,7 +237,7 @@ class ir(object):
         ir_blocs_all = []
         for l in bloc.lines:
             if c is None:
-                label = self.get_label(l)
+                label = self.get_instr_label(l)
                 c = irbloc(label, [], [])
                 ir_blocs_all.append(c)
             ir_bloc_cur, ir_blocs_extra = self.instr2ir(l)
@@ -290,9 +300,11 @@ class ir(object):
             self.blocs[irb.label] = irb
 
 
-    def get_label(self, instr):
-        l = self.symbol_pool.getby_offset_create(instr.offset)
-        return l
+    def get_instr_label(self, instr):
+        """Returns the label associated to an instruction
+        @instr: current instruction"""
+
+        return self.symbol_pool.getby_offset_create(instr.offset)
 
     def gen_label(self):
         # TODO: fix hardcoded offset