diff --git a/plugins/bap/plugins/bap_functions.py b/plugins/bap/plugins/bap_functions.py index ae512c3..bd5003f 100644 --- a/plugins/bap/plugins/bap_functions.py +++ b/plugins/bap/plugins/bap_functions.py @@ -10,12 +10,13 @@ import idaapi import idc +from heapq import heappush, heappop from bap.utils.run import BapIda class FunctionFinder(BapIda): def __init__(self): - super(FunctionFinder, self).__init__() + super(FunctionFinder, self).__init__(symbols=False) self.action = 'looking for function starts' self.syms = self.tmpfile('syms', mode='r') self.args += [ @@ -23,9 +24,18 @@ def __init__(self): '--dump', 'symbols:{0}'.format(self.syms.name) ] + # we can be a little bit more promiscuous since IDA will ignore + # function starts that occur in the middle of a function + if 'byteweight' in self.plugins and not \ + '--no-byteweight' in self.args: + self.args += [ + '--byteweight-threshold', '0.5', + '--byteweight-length', '4', + ] + class BAP_Functions(idaapi.plugin_t): - """Plugin to get functions from BAP and mark them in IDA.""" + """Uses BAP to find missed functions""" flags = idaapi.PLUGIN_FIX comment = "BAP Functions Plugin" @@ -40,15 +50,13 @@ def mark_functions(self): analysis.run() def add_starts(self, bap): - idaapi.refresh_idaview_anyway() + syms = [] for line in bap.syms: - line = line.strip() - if len(line) == 0: - continue - addr = int(line, 16) - end_addr = idaapi.BADADDR - idaapi.add_func(addr, end_addr) + heappush(syms, int(line, 16)) + for i in range(len(syms)): + idaapi.add_func(heappop(syms), idaapi.BADADDR) idc.Refresh() + idaapi.refresh_idaview_anyway() def init(self): """Initialize Plugin.""" diff --git a/plugins/bap/utils/run.py b/plugins/bap/utils/run.py index c9b9947..f225a26 100644 --- a/plugins/bap/utils/run.py +++ b/plugins/bap/utils/run.py @@ -23,14 +23,24 @@ class Bap(object): We will try to keep it clean from IDA specifics, so that later we can lift it to the bap-python library + + Attributes: + + DEBUG print executed commands and keep temporary files + args default arguments, inserted after `bap ` + plugins a list of available plugins """ DEBUG = False + args = [] + + plugins = [] + def __init__(self, bap, input_file): """Sandbox for the BAP process. - Each process is sandboxed, so that all intermediated data is + Each process is sandboxed, so that all intermediate data are stored in a temporary directory. instance variables: @@ -50,7 +60,7 @@ def __init__(self, bap, input_file): """ self.tmpdir = tempfile.mkdtemp(prefix="bap") - self.args = [bap, input_file] + self.args = [bap, input_file] + self.args self.proc = None self.fds = [] self.out = self.tmpfile("out") @@ -59,6 +69,9 @@ def __init__(self, bap, input_file): self.env = {'BAP_LOG_DIR': self.tmpdir} if self.DEBUG: self.env['BAP_DEBUG'] = 'yes' + if not Bap.plugins: + with os.popen(bap + ' --list-plugins') as out: + Bap.plugins = [e.split()[1] for e in out] def run(self): "starts BAP process"