Connect BasicBlock to collectors

Author: nirizr

idaplugin/rematch/collectors/vectors/__init__.py

@@ -5,8 +5,9 @@
 from .assembly_hash import AssemblyHashVector
 from .mnemonic_hash import MnemonicHashVector
 from .mnemonic_hist import MnemonicHistVector
+from .basicblock import BasicBlockVector
 
 
 __all__ = ["Vector", "InstructionHashVector", "IdentityHashVector",
            "NameHashVector", "AssemblyHashVector", "MnemonicHashVector",
-           "MnemonicHistVector"]
+           "MnemonicHistVector", "BasicBlockVector"]

idaplugin/rematch/collectors/vectors/basicblocks.py idaplugin/rematch/collectors/vectors/basicblock.py

@@ -1,18 +1,30 @@
 import ida_gdl
 import ida_funcs
 
+from . import vector
 
-class BasicBlocks(object):
-  def __init__(self, offset):
-    self.func = ida_funcs.get_func(offset)
+
+class BasicBlockVector(vector.Vector):
+  type = "basicblock"
+  type_version = 0
+
+  def __init__(self, *args, **kwargs):
+    super(BasicBlockVector, self).__init__(*args, **kwargs)
+
+    self.func = ida_funcs.get_func(self.offset)
 
     self.nodes = filter(self.node_contained, ida_gdl.FlowChart(self.func))
     self.node_ids = map(lambda n: n.id, self.nodes)
 
     self.seen_nodes = set()
 
+  def data(self, offset):
+    del offset
+
     # Assuming node #0 is the root node
-    print(self.add_node(self.nodes[0]))
+    serialized_bbs = self.add_node(self.nodes[0])
+    print(serialized_bbs)
+    return serialized_bbs
 
   def add_node(self, node):
     if node.id in self.seen_nodes:

idaplugin/rematch/instances/function.py

@@ -21,7 +21,8 @@ def __init__(self, *args, **kwargs):
                      collectors.vectors.IdentityHashVector,
                      collectors.vectors.AssemblyHashVector,
                      collectors.vectors.MnemonicHashVector,
-                     collectors.vectors.MnemonicHistVector}
+                     collectors.vectors.MnemonicHistVector,
+                     collectors.vectors.BasicBlockVector}
     self.annotations |= {collectors.annotations.AssemblyAnnotation}
 
   def size(self):