def set_cfg(line):
f = idaapi.get_func(line)
if f:
cfg = idaapi.FlowChart(f)
else:
cfg = []
return cfg
def readFunction(self, f, discard=True):
name = idc.GetFunctionName(f)
func = idaapi.get_func(f)
flow = idaapi.FlowChart(func)
size = func.endEA - func.startEA
if discard:
# Unnamed function, ignore it...
if name.startswith("sub_") or name.startswith(
"j_") or name.startswith("unknown"):
return False
# Already recognized runtime's function
flags = idc.GetFunctionFlags(f)
if flags & idc.FUNC_LIB or flags == -1:
return False
nodes = 0
edges = 0
points = 0
instructions = 0
mnems = []
dones = {}
for block in flow:
nodes += 1
indegree = 0
outdegree = 0
for succ_block in block.succs():
edges += 1
indegree += 1
if not dones.has_key(succ_block.id):
dones[succ_block] = 1
for x in list(
idautils.Heads(succ_block.startEA,
succ_block.endEA)):
instructions += 1
mnems.append(idc.GetMnem(x))
for pred_block in block.preds():
edges += 1
outdegree += 1
if not dones.has_key(succ_block.id):
dones[succ_block] = 1
for x in list(
idautils.Heads(succ_block.startEA,
succ_block.endEA)):
instructions += 1
mnems.append(idc.GetMnem(x))
if indegree > 0:
points += indegree
if outdegree > 0:
points += outdegree
if nodes > 1 and instructions > 5 and edges > 1:
#myexport_print("Exporter: Current function 0x%08x %s" % (f, name))
return (name, nodes, edges, points, size, instructions, mnems)
return False
def query_all_callback(self, threshold=0.8, minsize=3):
for ea in idautils.Functions():
pfn = idaapi.get_func(ea)
func_name = idaapi.get_func_name(ea)
if idaapi.FlowChart(pfn).size < minsize:
print(
"[BinaryAI] {} is skipped because basicblock size lower than minsize({})"
.format(func_name, minsize))
continue
funcs = self.query_function(ea)
if funcs is None:
print(
"[BinaryAI] {} is skipped because get function feature error"
.format(func_name, threshold))
continue
func = funcs[0]
if func['score'] < threshold:
print(
"[BinaryAI] {} is skipped because top1_score lower than threshold({})"
.format(func_name, threshold))
continue
idc.set_color(ea, idc.CIC_FUNC, 0xFFFFE1)
idc.set_func_flags(ea, idc.get_func_flags(ea) | 0x10000)
comment = SourceCodeViewer.source_code_comment(func_name, func)
idaapi.set_func_cmt(pfn, comment, 0)
def _match_with_check(self, ea, topk, funcset_ids):
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(ea,
topk=bai_config['topk'],
funcset_ids=funcset_ids)
except 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 analysis():
all_funcs = idautils.Functions()
overall_addr = dict()
for f in all_funcs:
f = idaapi.FlowChart(idaapi.get_func(f), flags=idaapi.FC_PREDS)
for block in f:
if block.startEA > idc.PrevHead(block.endEA):
continue
key = ''
# overall_addr.append(hex(block.startEA))
key += hex(block.startEA)
key += ','
key += hex(idc.PrevHead(block.endEA))
sus_addr = list()
successor = block.succs()
for addr in successor:
sus_addr.append(hex(addr.startEA))
overall_addr[key] = sus_addr
filename = idc.GetInputFile() + "_cfg"
with open(filename, 'w') as f:
json.dump(overall_addr, f)
def flow(key=None):
if key is None:
fn = ui.current.function()
if fn is None: raise LookupError, "function.bottom(%r):Not currently positioned within a function"% key
else:
fn = by(key)
fc = idaapi.FlowChart(f=fn, flags=idaapi.FC_PREDS)
return fc
