diff options
Diffstat (limited to 'miasm2/analysis/depgraph.py')
| -rw-r--r-- | miasm2/analysis/depgraph.py | 195 |
1 files changed, 168 insertions, 27 deletions
diff --git a/miasm2/analysis/depgraph.py b/miasm2/analysis/depgraph.py index 1b1d7ed4..a375af1b 100644 --- a/miasm2/analysis/depgraph.py +++ b/miasm2/analysis/depgraph.py @@ -1,6 +1,7 @@ """Provide dependency graph""" import itertools -from collections import namedtuple +from collections import deque +from UserDict import IterableUserDict try: import z3 @@ -23,7 +24,7 @@ class DependencyNode(object): line. """ - def __init__(self, label, element, line_nb, modifier=False): + def __init__(self, label, element, line_nb, step, modifier=False): """Create a dependency node with: @label: asm_label instance @element: Expr instance @@ -34,7 +35,10 @@ class DependencyNode(object): self._element = element self._line_nb = line_nb self._modifier = modifier - self._hash = hash((self._label, self._element, self._line_nb)) + self._step = step + self._nostep_repr = (self._label, self._line_nb, self._element) + self._hash = hash( + (self._label, self._element, self._line_nb, self._step)) def __hash__(self): return self._hash @@ -44,7 +48,8 @@ class DependencyNode(object): return False return (self.label == depnode.label and self.element == depnode.element and - self.line_nb == depnode.line_nb) + self.line_nb == depnode.line_nb and + self.step == depnode.step) def __cmp__(self, node): if not isinstance(node, self.__class__): @@ -62,6 +67,11 @@ class DependencyNode(object): return self.__str__() @property + def nostep_repr(self): + """Returns a representation of @self ignoring the step attribute""" + return self._nostep_repr + + @property def label(self): "Name of the current IRBlock" return self._label @@ -77,6 +87,11 @@ class DependencyNode(object): return self._line_nb @property + def step(self): + "Step of the current node" + return self._step + + @property def modifier(self): """Evaluating the current line involves a modification of tracked dependencies""" @@ -89,6 +104,47 @@ class DependencyNode(object): self._modifier = value +class CacheWrapper(IterableUserDict): + + """Wrapper class for cache dictionnary""" + + def __init__(self, dct=None): + """Create a CacheWrapper with value @dct.""" + IterableUserDict.__init__(self, dct) + self._nostep_cache = None + self._nostep_keys = None + + def __eq__(self, cache): + """Returns True if the nostep caches are equals""" + if self.nostep_keys != cache.nostep_keys: + return False + return self.nostep_cache == cache.nostep_cache + + @property + def nostep_keys(self): + """List of dictonnary keys without the step attribute. + The list is generated once when the method is called and not updated + afterward. + """ + if self._nostep_keys is None: + self._nostep_keys = set([key.nostep_repr for key in self.data]) + return self._nostep_keys + + @property + def nostep_cache(self): + """Dictionnary of DependencyNode and their dependencies, + without the step attribute. + The dictionnary is generated once when the method is called for the first time + and not updated afterward. + """ + if self._nostep_cache is None: + self._nostep_cache = {} + for (node, values) in self.data.iteritems(): + self._nostep_cache.setdefault(node.nostep_repr, set()).update( + set([val.nostep_repr for val in values])) + return self._nostep_cache + + class DependencyDict(object): """Internal structure for the DependencyGraph algorithm""" @@ -102,14 +158,13 @@ class DependencyDict(object): self._pending = set() # DepNode -> set(DepNode) - self._cache = {} + self._cache = CacheWrapper() def __eq__(self, depdict): if not isinstance(depdict, self.__class__): return False return (self._label == depdict.label and - self._cache == depdict.cache and - self._pending == depdict.pending) + self.cache == depdict.cache) def __cmp__(self, depdict): if not isinstance(depdict, self.__class__): @@ -198,7 +253,7 @@ class DependencyDict(object): parallels.append(self._get_modifiers_in_cache(depnode)) if parallels: - for parallel in itertools.product(*parallels): + for parallel in itertools.product(*[p for p in parallels if p]): out.update(parallel) return out @@ -207,8 +262,8 @@ class DependencyDict(object): """Remove intermediary states (non modifier depnodes) in the internal cache values""" - cache_out = {} - for depnode in self._cache.keys(): + cache_out = CacheWrapper() + for depnode in self._cache: cache_out[depnode] = self._get_modifiers_in_cache(depnode, force=True) self._cache = cache_out @@ -294,6 +349,74 @@ class DependencyDict(object): if key not in used_nodes: del self._cache[key] + def filter_unmodifier_loops(self, implicit, irdst): + """ + Remove unmodifier node creating dependency loops over + pending elements in cache. + """ + + previous_dict = None + # Get pending nodes of last time the label was handled + for hist_dict in reversed(self.history): + if hist_dict.label == self.label: + previous_dict = hist_dict + break + + if not previous_dict: + return + + nostep_pending = [node.nostep_repr for node in self.pending] + + to_remove = set() + for depnode in previous_dict.pending: + if (depnode.nostep_repr not in nostep_pending or + implicit and depnode.element == irdst): + continue + + to_remove.update(self._non_modifier_in_loop(depnode)) + + # Remove unused elements + for key in to_remove: + if depnode.nostep_repr == key.nostep_repr: + self._cache[depnode] = self._cache.get(key, set()).copy() + self.pending.discard(key) + self.pending.add(depnode) + if self._cache.has_key(key): + del self._cache[key] + + def _non_modifier_in_loop(self, depnode): + """ + Walk from @depnode until a node with the same nostep_repr is + encountered. + Returns a set of unmodifier nodes met in the path if no modifier was + found. + Returns set() if there exist a modifier node on the path. + """ + if not self.cache.has_key(depnode): + return set() + # Init + todo = set(self.cache[depnode]) + unmodifier_nodes = [] + + # Map + while todo: + node = todo.pop() + if node in unmodifier_nodes: + continue + if node.modifier: + return set() + unmodifier_nodes.append(node) + if not node in self._cache: + continue + if node.nostep_repr == depnode.nostep_repr: + unmodifier_nodes.append(node) + break + + for sub_node in self._cache[node]: + todo.add(sub_node) + + return unmodifier_nodes + class DependencyResult(object): """Container and methods for DependencyGraph results""" @@ -326,7 +449,8 @@ class DependencyResult(object): @property def unresolved(self): - return set(self._depdict.pending) + return set([node.nostep_repr for node in self._depdict.pending + if node.element != self._ira.IRDst]) @property def relevant_nodes(self): @@ -343,7 +467,7 @@ class DependencyResult(object): # Keep history order output = [] for label in [depdict.label for depdict in self.history]: - if label not in output and label in used_labels: + if label in used_labels: output.append(label) return output @@ -352,6 +476,14 @@ class DependencyResult(object): def input(self): return self._input_depnodes + @property + def has_loop(self): + """True if current depgraph represent a looping path""" + if self._has_loop is None: + self._has_loop = (len(self.relevant_labels) != + len(set(self.relevant_labels))) + return self._has_loop + def emul(self, ctx=None, step=False): """Symbolic execution of relevant nodes according to the history Return the values of input nodes' elements @@ -455,7 +587,7 @@ class FollowExpr(object): self.element = element @staticmethod - def to_depnodes(follow_exprs, label, line, modifier): + def to_depnodes(follow_exprs, label, line, modifier, counter): """Build a set of FollowExpr(DependencyNode) from the @follow_exprs set of FollowExpr""" dependencies = set() @@ -464,6 +596,7 @@ class FollowExpr(object): DependencyNode(label, follow_expr.element, line, + next(counter), modifier=modifier))) return dependencies @@ -501,6 +634,7 @@ class DependencyGraph(object): # Init self._ira = ira self._implicit = implicit + self._step_counter = itertools.count() # The IRA graph must be computed assert(hasattr(self._ira, 'g')) @@ -517,6 +651,11 @@ class DependencyGraph(object): self._cb_follow.append(self._follow_nocall) self._cb_follow.append(self._follow_label) + @property + def step_counter(self): + "Iteration counter" + return self._step_counter + @staticmethod def _follow_simp_expr(exprs): """Simplify expression so avoid tracking useless elements, @@ -531,13 +670,11 @@ class DependencyGraph(object): def _follow_label(exprs): """Do not follow labels""" follow = set() - unfollow = set() for expr in exprs: - if expr_is_label(expr): - unfollow.add(expr) - else: + if not expr_is_label(expr): follow.add(expr) - return follow, unfollow + + return follow, set() @staticmethod def _follow_mem_wrapper(exprs, mem_read): @@ -621,7 +758,8 @@ class DependencyGraph(object): ## Build output output = FollowExpr.to_depnodes(read, depnode.label, - depnode.line_nb - 1, modifier) + depnode.line_nb - 1, modifier, + self.step_counter) return output @@ -679,16 +817,16 @@ class DependencyGraph(object): # Init the work list done = {} - todo = [current_depdict] + todo = deque([current_depdict]) while todo: - depdict = todo.pop() + depdict = todo.popleft() # Update the dependencydict until fixed point is reached self._updateDependencyDict(depdict) # Clean irrelevant path - depdict.filter_used_nodes(depnodes) + depdict.filter_unmodifier_loops(self._implicit, self._ira.IRDst) # Avoid infinite loops label = depdict.label @@ -714,14 +852,16 @@ class DependencyGraph(object): 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) + irb_len, + next(self.step_counter), + modifier=False) ## Create links between DependencyDict for depnode_head in depdict.pending: ### Follow the head element in the parent new_depnode = DependencyNode(label, depnode_head.element, - irb_len) + irb_len, + next(self.step_counter)) ### The new node has to be computed in _updateDependencyDict new_depdict.cache[depnode_head] = set([new_depnode]) new_depdict.pending.add(new_depnode) @@ -729,6 +869,7 @@ class DependencyGraph(object): ### Handle implicit dependencies if self._implicit: new_depdict.cache[depnode_head].add(implicit_depnode) + new_depdict.pending.add(implicit_depnode) ## Manage the new element @@ -753,7 +894,8 @@ class DependencyGraph(object): # Init the algorithm input_depnodes = set() for element in elements: - input_depnodes.add(DependencyNode(label, element, line_nb)) + input_depnodes.add(DependencyNode(label, element, line_nb, + next(self.step_counter))) # Compute final depdicts depdicts = self._processInterBloc(input_depnodes, heads) @@ -791,4 +933,3 @@ class DependencyGraph(object): @heads: set of asm_label instances """ return self.get(label, elements, len(self._get_irs(label)), heads) - |