def check_ida():
if idaapi.IDA_SDK_VERSION < 700:
BinaryAILog.log(BinaryAILog.ERROR, "Need IDA >= 7.0")
return False
if not idaapi.init_hexrays_plugin():
BinaryAILog.log(BinaryAILog.ERROR,
"Hex-Rays decompiler does not exists")
return False
return True
def init(self):
if not idaapi.is_idaq():
BinaryAILog.log(BinaryAILog.INFO,
"Plugin should be loaded in idaq mode")
return idaapi.PLUGIN_SKIP
if check_ida():
bai_mgr = BinaryAIManager()
ui_mgr = UIManager(BinaryAIIDAPlugin.wanted_name, bai_mgr)
if ui_mgr.register_actions():
return idaapi.PLUGIN_KEEP
else:
BinaryAILog.log(BinaryAILog.ERROR, "Register actions failed")
return idaapi.PLUGIN_SKIP
def upload_funcs(self, funcs):
i, succ, skip, fail = 0, 0, 0, 0
def stop():
idaapi.hide_wait_box()
BinaryAILog.summary(succ, skip, fail, "uploaded")
funcs_len = len(funcs)
idaapi.show_wait_box("Uploading... (0/{})".format(funcs_len))
for ea in funcs:
i += 1
idaapi.replace_wait_box("Uploading... ({}/{})".format(
i, funcs_len))
if idaapi.user_cancelled():
stop()
return
# < minsize
pfn = idaapi.get_func(ea)
if idaapi.FlowChart(pfn).size < bai_config['minsize']:
skip += 1
continue
# try upload
func_id = None
try:
func_id = self.mgr.upload(ea)
except idaapi.DecompilationFailure as e:
BinaryAILog.fail(idaapi.get_func_name(ea), str(e))
fail += 1
continue
except BinaryAIException as e:
stop()
BinaryAILog.fatal(e)
# fail
if not func_id:
fail += 1
continue
succ += 1
stop()
def _match_with_check(self, ea):
fail, skip, succ = -1, 0, 1
# < minsize
pfn = idaapi.get_func(ea)
if idaapi.FlowChart(pfn).size < bai_config['minsize']:
return skip
# do match
try:
targets = self.mgr.retrieve_by_feature(ea, topk=1)
except idaapi.DecompilationFailure as e:
BinaryAILog.fail(idaapi.get_func_name(ea), str(e))
return fail
except BinaryAIException as e:
idaapi.hide_wait_box()
BinaryAILog.fatal(e)
if targets is None:
return fail
if targets[0]['score'] < bai_config['threshold'] or \
not bai_mark.apply_bai_high_score(
ea,
targets[0]['function']['name'],
targets[0]['score']):
return skip
return succ
def upload(self, ea):
try:
func_id = self.mgr.upload(ea)
except idaapi.DecompilationFailure as e:
BinaryAILog.fail(idaapi.get_func_name(ea), str(e))
except BinaryAIException as e:
BinaryAILog.fatal(e)
if func_id:
BinaryAILog.success(idaapi.get_func_name(ea), func_id, "uploaded")
def retrieve(self, ea, cview):
if not self.check_before_use():
return
func_name = idaapi.get_func_name(ea)
try:
targets = self.mgr.retrieve_by_feature(ea, bai_config['topk'])
except idaapi.DecompilationFailure as e:
BinaryAILog.fail(idaapi.get_func_name(ea), str(e))
except BinaryAIException as e:
BinaryAILog.fatal(e)
if targets is None:
BinaryAILog.skip(func_name, "get function feature error")
return
if len(targets) == 0:
idaapi.warning("No similar function found!")
return
cview.set_user_data(ea, targets)
def stop():
idaapi.hide_wait_box()
BinaryAILog.summary(succ, skip, 0, "reverted")
def stop():
idaapi.hide_wait_box()
BinaryAILog.summary(succ, skip, fail, "uploaded")