Convert vectors from class-based to object-based

So init methods could be useful. Needed this in the basicblocks vector

Author: nirizr

idaplugin/rematch/collectors/__init__.py

@@ -9,7 +9,8 @@
 def collect(collectors, offset, instance_id=None):
   for collector in collectors:
     try:
-      r = collector.collect(offset, instance_id)
+      collector_obj = collector(offset, instance_id)
+      r = collector_obj.collect()
       if r:
         yield r
     except UnicodeDecodeError:

idaplugin/rematch/collectors/vectors/assembly_hash.py

@@ -10,13 +10,12 @@ class AssemblyHashVector(vector.Vector):
   type = 'assembly_hash'
   type_version = 0
 
-  @classmethod
-  def data(cls, offset):
-    if cls.inst_count(offset) < 3:
+  def data(self):
+    if self.inst_count() < 3:
       return None
 
     md5 = hashlib.md5()
-    for ea in idautils.FuncItems(offset):
+    for ea in idautils.FuncItems(self.offset):
       asm_line = idc.GetDisasmEx(ea, idc.GENDSM_MULTI_LINE)
       if ';' in asm_line:
         asm_line = asm_line[:asm_line.find(';')]

idaplugin/rematch/collectors/vectors/basicblock.py

@@ -18,8 +18,7 @@ def __init__(self, *args, **kwargs):
 
     self.seen_nodes = set()
 
-  def data(self, offset):
-    del offset
+  def data(self):
 
     # Assuming node #0 is the root node
     serialized_bbs = self.add_node(self.nodes[0])
@@ -48,7 +47,7 @@ def token(node):
     # alternatives: offset from start of function
     return node.endEA - node.startEA
 
-  sort_ket = token
+  sort_key = token
 
   def node_contained(self, node):
     # make sure only nodes inside the function are accounted for

idaplugin/rematch/collectors/vectors/identity_hash.py

@@ -21,17 +21,16 @@ def _cycle(h, b):
     h &= 0xffffffffffffffff
     return h
 
-  @classmethod
-  def data(cls, offset):
-    if cls.inst_count(offset) < 3:
+  def data(self):
+    if self.inst_count() < 3:
       return None
 
-    h = cls.keleven
-    for ea in idautils.FuncItems(offset):
-      h = cls._cycle(h, idc.Byte(ea))
+    h = self.keleven
+    for ea in idautils.FuncItems(self.offset):
+      h = self._cycle(h, idc.Byte(ea))
       # skip additional bytes of any instruction that contains an offset in it
       if idautils.CodeRefsFrom(ea, False) or idautils.DataRefsFrom(ea):
         continue
       for i in range(ea + 1, ea + idc.ItemSize(ea)):
-        h = cls._cycle(h, idc.Byte(i))
+        h = self._cycle(h, idc.Byte(i))
     return h

idaplugin/rematch/collectors/vectors/instruction_hash.py

@@ -21,15 +21,14 @@ def _cycle(h, b):
     h &= 0xffffffffffffffff
     return h
 
-  @classmethod
-  def data(cls, offset):
-    if cls.inst_count(offset) < 3:
+  def data(self):
+    if self.inst_count() < 3:
       return None
 
-    h = cls.keleven
-    for ea in idautils.FuncItems(offset):
-      h = cls._cycle(h, idc.Byte(ea))
+    h = self.keleven
+    for ea in idautils.FuncItems(self.offset):
+      h = self._cycle(h, idc.Byte(ea))
       # go over all additional bytes of any instruction
       for i in range(ea + 1, ea + idc.ItemSize(ea)):
-        h = cls._cycle(h, idc.Byte(i))
+        h = self._cycle(h, idc.Byte(i))
     return h

idaplugin/rematch/collectors/vectors/mnemonic_hash.py

@@ -10,13 +10,12 @@ class MnemonicHashVector(vector.Vector):
   type = 'mnemonic_hash'
   type_version = 0
 
-  @classmethod
-  def data(cls, offset):
-    if cls.inst_count(offset) < 3:
+  def data(self):
+    if self.inst_count() < 3:
       return None
 
     md5 = hashlib.md5()
-    for ea in idautils.FuncItems(offset):
+    for ea in idautils.FuncItems(self.offset):
       mnem_line = idc.GetMnem(ea)
       mnem_line = mnem_line.strip()
       mnem_line = mnem_line.lower()

idaplugin/rematch/collectors/vectors/mnemonic_hist.py

@@ -10,11 +10,10 @@ class MnemonicHistVector(vector.Vector):
   type = 'mnemonic_hist'
   type_version = 0
 
-  @classmethod
-  def data(cls, offset):
+  def data(self):
     instruction_hist = defaultdict(int)
 
-    for ea in idautils.FuncItems(offset):
+    for ea in idautils.FuncItems(self.offset):
       mnem_line = idc.GetMnem(ea)
       mnem_line = mnem_line.lower()
       instruction_hist[mnem_line] += 1

idaplugin/rematch/collectors/vectors/name_hash.py

@@ -10,9 +10,8 @@ class NameHashVector(vector.Vector):
   type = 'name_hash'
   type_version = 0
 
-  @classmethod
-  def data(cls, offset):
-    name = idc.Name(offset)
+  def data(self):
+    name = idc.Name(self.offset)
     if ida_name.is_uname(name):
       return None
 

idaplugin/rematch/collectors/vectors/vector.py

@@ -4,16 +4,18 @@
 
 
 class Vector(object):
-  @classmethod
-  def collect(cls, offset, instance_id=None):
-    data = cls.data(offset)
+  def __init__(self, offset, instance_id=None):
+    self.offset = offset
+    self.instance_id = instance_id
+
+  def collect(self):
+    data = self.data()
     if not data:
       return None
 
     data = json.dumps(data)
-    return {"instance": instance_id, "type": cls.type,
-            "type_version": cls.type_version, "data": data}
+    return {"instance": self.instance_id, "type": self.type,
+            "type_version": self.type_version, "data": data}
 
-  @staticmethod
-  def inst_count(offset):
-    return len(list(idautils.FuncItems(offset)))
+  def inst_count(self):
+    return len(list(idautils.FuncItems(self.offset)))

server/collab/matchers/__init__.py

@@ -8,14 +8,16 @@
 from .name_hash import NameHashMatcher
 from .mnemonic_euclidean import MnemonicEuclideanMatcher
 from .dictionary_matcher import DictionaryMatcher
+from .basicblock import BasicBlockMatcher
 
 
 matchers_list = [InstructionHashMatcher, IdentityHashMatcher, NameHashMatcher,
                  AssemblyHashMatcher, MnemonicHashMatcher,
-                 MnemonicEuclideanMatcher]
+                 MnemonicEuclideanMatcher, BasicBlockMatcher]
 
 
 __all__ = ['Matcher', 'HashMatcher', 'EuclideanDictionaryMatcher',
            'InstructionHashMatcher', 'IdentityHashMatcher',
            'AssemblyHashMatcher', 'MnemonicHashMatcher', 'NameHashMatcher',
-           'MnemonicEuclideanMatcher', 'DictionaryMatcher', 'matchers_list']
+           'MnemonicEuclideanMatcher', 'DictionaryMatcher',
+           'BasicBlockMatcher', 'matchers_list']

server/collab/matchers/basicblock.py

@@ -0,0 +1,8 @@
+from . import matcher
+
+
+class BasicBlockMatcher(matcher.Matcher):
+  vector_type = 'basicblock'
+  match_type = 'basicblock'
+  matcher_name = "Basic Block"
+  matcher_description = ("TODO")