diff --git a/idaplugin/rematch/collectors/annotations/__init__.py b/idaplugin/rematch/collectors/annotations/__init__.py index d9a50f99b..c7341786b 100644 --- a/idaplugin/rematch/collectors/annotations/__init__.py +++ b/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"] diff --git a/idaplugin/rematch/collectors/annotations/comment.py b/idaplugin/rematch/collectors/annotations/comment.py new file mode 100644 index 000000000..85063a600 --- /dev/null +++ b/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) diff --git a/idaplugin/rematch/collectors/annotations/positional.py b/idaplugin/rematch/collectors/annotations/positional.py new file mode 100644 index 000000000..363af96cb --- /dev/null +++ b/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 diff --git a/idaplugin/rematch/instances/function.py b/idaplugin/rematch/instances/function.py index f00603a3a..336ad67b3 100644 --- a/idaplugin/rematch/instances/function.py +++ b/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