From 289fd86c1ee9e0083ae832cfe1b9beb839d25a53 Mon Sep 17 00:00:00 2001
From: Ivan Gotovchits <ivg@ieee.org>
Date: Thu, 29 Jun 2017 14:49:53 -0400
Subject: [PATCH] fixes and enhances function recover plugin

The function recovery plugin was broken for some amount of time, as we
passed symbol file to bap, effectively disabling BAP own function
finding capabilities, so bap always returned the same amount of function
starts.

This commit also enhances the recovery plugin, by lowering the recovery
threshold. Although it will lead to more false positives it is
acceptable since IDA will ignore function starts that occur in a body of
an discovered function, that was already discovered and
disassembled. Since most of the functions have their bodies after
starts, we added sorting of the addresses, so that function starts are
added to the system in the ascending order.

This commit also adds two new attributes to the BapIda class. The first
one is called [args] and is a list of default arguments, shared by all
instances. It is useful to adapt the behavior of the plugins to your
needs, and to perform fast prototyping and debugging.

The second attribute is called [plugins] and it contains a list of
available BAP plugins. This is useful, when an IDA plugin needs to adapt
its behavior based on the presence of particular BAP plugins. For
example, the function start identification plugin, is allowed to tackle
with the byteweight parameters only if it is installed.
---
 plugins/bap/plugins/bap_functions.py | 26 +++++++++++++++++---------
 plugins/bap/utils/run.py             | 17 +++++++++++++++--
 2 files changed, 32 insertions(+), 11 deletions(-)

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 <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"