diff --git a/plugins/codatify/codatify.py b/plugins/codatify/codatify.py index 7294bdb..dba547c 100755 --- a/plugins/codatify/codatify.py +++ b/plugins/codatify/codatify.py @@ -256,10 +256,17 @@ def stringify(self): print ("Looking for possible strings starting at: %s:0x%X..." % (idc.SegName(ea), ea)), + isbreak = False for s in idautils.Strings(): + if isbreak is True: + break if s.ea > ea: - if not idc.isASCII(idc.GetFlags(s.ea)) and idc.MakeStr(s.ea, idc.BADADDR): - n += 1 + try: + if not idc.isASCII(idc.GetFlags(s.ea)) and idc.MakeStr(s.ea, idc.BADADDR): + n += 1 + except: + isbreak = True + pass print "created %d new ASCII strings" % n @@ -354,6 +361,37 @@ def codeify(self, ea=idc.BADADDR): print "Created %d new functions and %d new code blocks\n" % (func_count, code_count) +label_fixup_code = "Fixup Code" +label_fixup_data = "Fixup Data" +action_fixup_code = "fixup:code" +action_fixup_data = "fixup:data" +options_codatify = "Options/Codatify/" + +class CodatifyHandler(idaapi.action_handler_t): + def __init__(self, action): + idaapi.action_handler_t.__init__(self) + self.action = action + + def activate(self, ctx): + if self.action is 0: + self.fix_code() + elif self.action is 1: + self.fix_data() + return 1 + + def fix_code(self): + cd = Codatify() + cd.codeify() + + def fix_data(self): + cd = Codatify() + cd.stringify() + cd.datify() + cd.pointify() + StructFinder().parse_function_tables() + + def update(self, ctx): + return idaapi.AST_ENABLE_FOR_FORM if ctx.form_type == idaapi.BWN_DISASM else idaapi.AST_DISABLE_FOR_FORM class codatify_t(idaapi.plugin_t): flags = 0 @@ -363,28 +401,34 @@ class codatify_t(idaapi.plugin_t): wanted_hotkey = "" def init(self): - self.menu_context = idaapi.add_menu_item("Options/", "Fixup code", "", 0, self.fix_code, (None,)) - self.menu_context = idaapi.add_menu_item("Options/", "Fixup data", "", 0, self.fix_data, (None,)) + if idaapi.IDA_SDK_VERSION >= 700: + if idaapi.register_action(idaapi.action_desc_t(action_fixup_code,label_fixup_code,CodatifyHandler(0))) is None: + print "Fail register action fixup code" + if idaapi.register_action(idaapi.action_desc_t(action_fixup_data,label_fixup_data,CodatifyHandler(1))) is None: + print "Fail register action fixup data" + + if idaapi.attach_action_to_menu(options_codatify, action_fixup_code, idaapi.SETMENU_APP) is None: + print "Failed attach action" + if idaapi.attach_action_to_menu(options_codatify, action_fixup_data, idaapi.SETMENU_APP) is None: + print "Failed attach action" + else: + self.menu_context = idaapi.add_menu_item(options_codatify, label_fixup_code, "", 0, CodatifyHandler(0), (None,)) + self.menu_context = idaapi.add_menu_item(options_codatify, label_fixup_data, "", 0, CodatifyHandler(1), (None,)) return idaapi.PLUGIN_KEEP def term(self): - idaapi.del_menu_item(self.menu_context) + if idaapi.IDA_SDK_VERSION >= 700: + if idaapi.unregister_action(action_fixup_code) is None: + print "Failed to unregister action." + if idaapi.unregister_action(action_fixup_data) is None: + print "Failed to unregister action." + else: + idaapi.del_menu_item(self.menu_context) return None def run(self, arg): pass - def fix_code(self, arg): - cd = Codatify() - cd.codeify() - - def fix_data(self, arg): - cd = Codatify() - cd.stringify() - cd.datify() - cd.pointify() - StructFinder().parse_function_tables() - def PLUGIN_ENTRY(): return codatify_t()