Add positional annotation base and a single comment annotation

idaplugin/rematch/collectors/annotations/__init__.py

@@ -2,7 +2,10 @@
 from .name import NameAnnotation
 from .assembly import AssemblyAnnotation
 from .prototype import PrototypeAnnotation
+from .positional import PositionalAnnotation
+from comment import RegularCommentAnnotation
 
 
 __all__ = ["Annotation", "NameAnnotation", "AssemblyAnnotation",
-           "PrototypeAnnotation"]
+           "PrototypeAnnotation", "PositionalAnnotation",
+           "RegularCommentAnnotation"]

idaplugin/rematch/collectors/annotations/comment.py

@@ -0,0 +1,23 @@
+import ida_funcs
+import ida_bytes
+import idautils
+
+from . import annotation
+
+
+class CommentAnnotation(annotation.Annotation):
+  def data(self):
+    func = ida_funcs.get_func(self.offset)
+
+    comments = {ea - self.offset: self.get_comment(ea)
+                  for ea in idautils.Heads(func.startEA, func.endEA)}
+    return comments
+
+  def get_comment(self, ea):
+    raise NotImplementedError("get_comment method not implemented")
+
+
+class RegularCommentAnnotation(CommentAnnotation):
+  @staticmethod
+  def get_comment(ea):
+    return ida_bytes.get_cmt(ea, 0)

idaplugin/rematch/collectors/annotations/positional.py

@@ -0,0 +1,47 @@
+import ida_gdl
+import ida_funcs
+import ida_lines
+import idautils
+
+from . import annotation
+
+
+class PositionalAnnotation(annotation.Annotation):
+  type = 'positional'
+
+  @staticmethod
+  def data(offset):
+    func = ida_funcs.get_func(offset)
+
+    def clean(asm):
+      asm = ida_lines.tag_remove(asm)
+      asm = asm.split(';', 1)[0]
+      asm.strip()
+      return asm
+
+    # make sure only nodes inside the function are accounted for
+    # this solves cascaded functions (when multiple functions share same ends)
+    def node_contained(node):
+      return (ida_funcs.func_contains(func, node.startEA) and
+              ida_funcs.func_contains(func, node.endEA - 1))
+    nodes = filter(node_contained, ida_gdl.FlowChart(func))
+    node_ids = map(lambda n: n.id, nodes)
+
+    nodes_data = []
+    for node in nodes:
+      assembly = [clean(ida_lines.generate_disasm_line(ea))
+                    for ea in idautils.Heads(node.startEA, node.endEA)]
+      successive_nodes = [succ.id
+                            for succ in node.succs()
+                            if succ.id in node_ids]
+      serialized_node = {'id': node.id, 'start': node.startEA,
+                         'end': node.endEA, 'successive': successive_nodes,
+                         'assembly': assembly}
+      nodes_data.append(serialized_node)
+
+    return nodes_data
+
+  @staticmethod
+  def apply(offset, data):
+    del offset
+    del data

idaplugin/rematch/instances/function.py

@@ -23,7 +23,9 @@ def __init__(self, *args, **kwargs):
                      collectors.vectors.MnemonicHashVector,
                      collectors.vectors.MnemonicHistVector,
                      collectors.vectors.BasicBlockSizeHistVector}
-    self.annotations |= {collectors.annotations.AssemblyAnnotation}
+    self.annotations |= {collectors.annotations.AssemblyAnnotation,
+                         collectors.annotations.PositionalAnnotation,
+                         collectors.annotations.RegularCommentAnnotation}
 
   def size(self):
     """return the overall size of function by adding sizes of all indevidual