Source Code for Module envi.bytesig

  1   
  2  """ 
  3  A byte and mask based decision engine for creating byte 
  4  sequences (and potential comparison masks) for general purpose 
  5  signature matching. 
  6   
  7  Currently used by vivisect function entry sig db and others. 
  8  """ 
  9   
 10 -class SignatureTree: 
 11      """ 
 12      A byte based decision tree which uses all the RAMs but is 
 13      really fast.... 
 14   
 15      Signatures consist of a byte sequence and an optional mask 
 16      sequence.  If present each mask byte is used to logical and 
 17      the byte being compared before comparison.  This allows the 
 18      creation of signatures which have parts of the sig generalized. 
 19   
 20      FIXME allow sigs to have a reliability rating 
 21      FIXME allow sig nodes to store depth and truncate the tree early (and then mask the rest) 
 22      """ 
 23   
 24 -    def __init__(self): 
 25          self.basenode = (0, [], [None] * 256) 
 26          self.sigs = {} # track duplicates 
 27   
 28 -    def _addChoice(self, siginfo, node): 
 29   
 30          todo = [(node, siginfo),] 
 31   
 32          # Recursion is for the weak minded. 
 33          while len(todo): 
 34   
 35              node, siginfo = todo.pop() 
 36              depth, sigs, choices = node 
 37              bytes, masks, o = siginfo 
 38              siglen = len(sigs) 
 39   
 40              sigs.append(siginfo) 
 41              if siglen == 0: 
 42                  pass # we just don't want the "else" here, if we're the only 
 43                  # one on this node, just let it ride. 
 44   
 45              elif siglen == 1: 
 46                  # If it has one already, we *both* need to add another level 
 47                  # (because if it is the only one, it thought it was last choice) 
 48                  for s in sigs: 
 49                      chval = s[0][depth] # get the choice byte from siginfo 
 50                      nnode = self._getNode(depth, choices, chval) 
 51                      todo.append((nnode, s)) 
 52   
 53              else: # This is already a choice node, keep on choosing... 
 54                  chval = bytes[depth] 
 55                  nnode = self._getNode(depth, choices, chval) 
 56                  todo.append((nnode, siginfo)) 
 57   
 58 -    def _getNode(self, depth, choices, choice): 
 59          # Chose, (and or initialize) a sub node 
 60          nnode = choices[choice] 
 61          if nnode == None: 
 62              nnode = (depth+1, [], [None] * 256) 
 63              choices[choice] = nnode 
 64          return nnode 
 65   
 66 -    def addSignature(self, bytes, masks=None, val=None): 
 67          """ 
 68          Add a signature to the search tree.  If masks goes unspecified, it will be 
 69          assumed to be all ones (\\xff * len(bytes)). 
 70   
 71          Additionally, you may specify "val" as the object to get back with 
 72          getSignature(). 
 73          """ 
 74          # FIXME perhaps make masks None on all ff's 
 75          if masks == None: 
 76              masks = "\xff" * len(bytes) 
 77   
 78          if val == None: 
 79              val = True 
 80   
 81          # Detect and skip duplicate additions... 
 82          bytekey = bytes + masks 
 83          if self.sigs.get(bytekey) != None: 
 84              return 
 85   
 86          self.sigs[bytekey] = True 
 87   
 88          byteord = [ord(c) for c in bytes] 
 89          maskord = [ord(c) for c in masks] 
 90   
 91          siginfo = (byteord, maskord, val) 
 92          self._addChoice(siginfo, self.basenode) 
 93   
 94 -    def isSignature(self, bytes, offset=0): 
 95          return self.getSignature(bytes, offset=offset) != None 
 96   
 97 -    def getSignature(self, bytes, offset=0): 
 98   
 99          node = self.basenode 
100          while True: 
101              depth, sigs, choices = node 
102              # Once we get down to one sig, there are no more branches, 
103              # just check the byte sequence. 
104              if len(sigs) == 1: 
105                  sbytes, smasks, sobj = sigs[0] 
106                  for i in xrange(depth, len(sbytes)): 
107                      realoff = offset + i 
108                      masked = ord(bytes[realoff]) & smasks[i] 
109                      if masked != sbytes[i]: 
110                          return None 
111                  return sobj 
112   
113              # There are still more choices, keep branching. 
114              node = None # Lets go find a new one 
115              for sig in sigs: 
116                  sbytes, smasks, sobj = sig 
117                  masked = ord(bytes[offset+depth]) & smasks[depth] 
118                  if sbytes[depth] == masked: # We have a winner! 
119                      # FIXME find the *best* winner! (because of masking) 
120                      node = choices[masked] 
121                      break 
122   
123              # We failed to make our next choice 
124              if node == None: 
125                  return None 
126