Merge pull request #29 from ivg/fixes-function-finding-plugin

fixes and enhances function recover plugin

Author: ivg

plugins/bap/plugins/bap_functions.py

@@ -10,22 +10,32 @@
 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 += [
             '--print-symbol-format', 'addr',
             '--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."""

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 <input>`
+       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"