Source Code for Module Libs.libcontrolflow

  1  #!/usr/bin/env python 
  2   
  3  """ 
  4  (c) Immunity, Inc. 2004-2007 
  5   
  6   
  7  U{Immunity Inc.<http://www.immunityinc.com>} 
  8   
  9   
 10  """ 
 11   
 12  __VERSION__ = '1.0' 
 13   
 14  ############################################################################# 
 15 -class DominatorTree: 
 16 -    def __init__(self, imm, addr, blocks = False, recursion = False): 
 17          """ 
 18          This class takes a function start address and calculate all Dominator Tree related tables: 
 19          - Predecessors 
 20          - Iterated Predecessors 
 21          - Dominators 
 22          - Immediate Dominators 
 23          - Post Dominators 
 24          - Immediate Post Dominators 
 25   
 26          @type  imm: Debbuger OBJECT 
 27          @param imm: Debbuger 
 28   
 29          @type  addr: DWORD 
 30          @param addr: function start address 
 31   
 32          @type  blocks: DICTIONARY|False 
 33          @param blocks: Optionally you can provide a dictionary with the node address as key and a list of edges (mainly for testing purposes). 
 34          """ 
 35           
 36          self.address = addr 
 37          self.imm = imm 
 38          self.blocks = {} 
 39          self.predecessors = {} 
 40          self.iterativepredecessors = {} 
 41          self.dominators = {} 
 42          self.immediatedominators = {} 
 43          self.postdominators = {} 
 44          self.immediatepostdominators = {} 
 45           
 46          if blocks: 
 47              self.blocks = blocks 
 48          else: 
 49              self.Initializate() 
 50           
 51          self.CalculatePredecessors() 
 52          self.CalculateDominators() 
 53          self.CalculateImmediateDominators() 
 54          if not recursion: 
 55              self.CalculatePostAndImmediatePostDominators() 
 56              self.CalculateIterativePredecessors() 
 57   
 58 -    def Initializate(self): 
 59          func = self.imm.getFunction(self.address) 
 60          blocks = func.getBasicBlocks() 
 61           
 62          for block in blocks: 
 63              edges = block.getEdges() 
 64              start = block.getStart() 
 65              self.blocks[start] = edges 
 66           
 67       
 68 -    def CalculatePredecessors(self): 
 69          for start,edges in self.blocks.iteritems(): 
 70              #support an unknown quantity of edges (for inverse CFG processing) 
 71              for edge in edges: 
 72                  if edge: 
 73                      if edge not in self.predecessors.keys(): 
 74                          self.predecessors[edge] = [] 
 75                      self.predecessors[edge].append(start) 
 76       
 77 -    def CalculateIterativePredecessors(self): 
 78          for start in self.blocks: 
 79              self.iterativepredecessors[start] = [] 
 80              if start in self.predecessors.keys(): 
 81                  self.__iterative_predecessors_helper(start, start) 
 82           
 83 -    def __iterative_predecessors_helper(self, base, newbase): 
 84          for pred in self.predecessors[newbase]: 
 85              if pred: 
 86                  if newbase in self.dominators[pred]: 
 87                      #this is a loop 
 88                      continue 
 89                  if pred not in self.iterativepredecessors[base]: 
 90                      self.iterativepredecessors[base].append(pred) 
 91                  if pred in self.predecessors.keys(): 
 92                      self.__iterative_predecessors_helper(base, pred) 
 93           
 94 -    def CalculateDominators(self): 
 95          """  
 96          Based in algorithm from "Advanced COMPILER DESIGN IMPLEMENTATION" 
 97          """ 
 98           
 99          start = self.address 
100          change = True 
101          Domin = {} 
102          Domin[start] = [ start ] 
103          for n in self.blocks: 
104              if n != start: 
105                  if n in self.predecessors.keys(): 
106                      Domin[n] = self.blocks.keys() 
107                  else: 
108                      #a node without predecessors it's just dead code 
109                      Domin[n] = [ n ] 
110   
111          for n in Domin: 
112              tmp = Domin[n] 
113              tmp.sort() 
114              Domin[n] = tmp 
115   
116          while change: 
117              change = False 
118              for n in self.blocks: 
119                  if n != start and n in self.predecessors.keys(): 
120                      T = self.blocks.keys() 
121                      for p in self.predecessors[n]: 
122                          #intersect Domin(p) with tmp 
123                          intersect = [] 
124                          for d in Domin[p]: 
125                              if d in T and d not in intersect: 
126                                  intersect.append(d) 
127                          T = intersect 
128   
129                      #D = T U n 
130                      D = intersect 
131                      if n not in D: 
132                          D.append(n) 
133                       
134                      D.sort() 
135                      if D != Domin[n]: 
136                          change = True 
137                          Domin[n] = D 
138           
139          self.dominators = Domin 
140   
141 -    def CalculateImmediateDominators(self): 
142          for node in self.blocks: 
143              idom = self.dominators[node][:] 
144              #idom(node) != node 
145              idom.remove(node) 
146              for dom in self.dominators[node]: 
147                  if dom != node: 
148                      for sec_dom in self.dominators[dom]: 
149                          if sec_dom != dom and sec_dom in idom: 
150                              idom.remove(sec_dom) 
151              self.immediatedominators[node] = idom 
152       
153 -    def CalculatePostAndImmediatePostDominators(self): 
154          invertedCFG = self.predecessors 
155          invertedCFG[self.address] = [ 0 ] 
156           
157          newstart = invertedCFG.keys() 
158          for edges in invertedCFG.values(): 
159              for edge in edges: 
160                  if edge in newstart: 
161                      newstart.remove(edge) 
162   
163          for onestart in newstart: 
164              dom = DominatorTree(self.imm, onestart, blocks=invertedCFG, recursion=True) 
165              self.postdominators[onestart]=dom.dominators 
166              self.immediatepostdominators[onestart]=dom.immediatedominators 
167   
168 -    def getDominators(self): 
169          return self.dominators 
170   
171 -    def getImmediateDominators(self): 
172          return self.immediatedominators 
173   
174 -    def getPostDominators(self): 
175          return self.postdominators 
176   
177 -    def getImmediatePostDominators(self): 
178          return self.immediatepostdominators 
179   
180 -    def getPredecessors(self): 
181          return self.predecessors 
182   
183 -    def getIteratedPredecessors(self): 
184          return self.iterativepredecessors 
185   
186 -    def getControlFlowGraph(self): 
187          return self.blocks 
188   
189   
190 -class ControlFlowAnalysis: 
191 -    def __init__(self, imm, address, domtree=False): 
192          """ 
193          @type  imm: Debbuger OBJECT 
194          @param imm: Debbuger 
195   
196          @type  address: DWORD 
197          @param address: function start address 
198   
199          @type  domtree: OBJECT|False 
200          @param domtree: Optionally you can provide a DominatorTree instance (mainly for testing purposes). 
201          """ 
202           
203          self.imm = imm 
204          self.address = address 
205          self.loops = [] 
206           
207          if domtree: 
208              self.domtree = domtree 
209          else: 
210              self.domtree = DominatorTree(self.imm, self.address) 
211       
212 -    def findNaturalLoops(self): 
213          """ 
214          This function finds Natural Loops inside a function, using the information provided by dominator tree class. 
215           
216          @rtype: LIST 
217          @return: A list of loops, each with this structure: 
218            [ start, end, nodes ] 
219            start: address of node receiving the back edge 
220            end: address of node which has the back edge 
221            node: list of node's addresses involved in this loop 
222          """ 
223           
224          for start,edges in self.domtree.blocks.items(): 
225              for edge in edges: 
226                  if edge and edge in self.domtree.dominators[start]: 
227                      loopNodes = [] 
228                      for pred in self.domtree.iterativepredecessors[start]: 
229                          if pred not in self.domtree.iterativepredecessors[edge]: 
230                              loopNodes.append(pred) 
231                      loopNodes.append(start) 
232                      self.loops.append([edge,start,loopNodes]) 
233           
234          return self.loops 
235