Add unit tests for address rebasing.

5 files changed, 154 insertions, 2 deletions

diff --git a/suite/api_test.py b/suite/api_test.py
index 5a1822aa..d30bcfbe 100644
--- a/suite/api_test.py
+++ b/suite/api_test.py
@@ -1,12 +1,24 @@
 import unittest
 import platform
 import os
+import sys
+import pickle
+import zipfile
+import difflib
+from collections import Counter
+
 from binaryninja.binaryview import BinaryView, BinaryViewType
 from binaryninja.settings import Settings, SettingsScope
 from binaryninja.metadata import Metadata
 from binaryninja.demangle import demangle_gnu3, get_qualified_name
 from binaryninja.architecture import Architecture
 
+api_suite_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../../", "api", "suite")
+sys.path.append(api_suite_path)
+import testcommon
+
+global verbose
+verbose = False
 
 class SettingsAPI(unittest.TestCase):
 	@classmethod
@@ -172,6 +184,107 @@ class SettingsAPI(unittest.TestCase):
 		assert mapped_view.segments[0].start == 0x500000, "test_load_settings failed"
 		assert len(mapped_view) == 4, "test_load_settings failed"
 
+
+class RebaseAPI(unittest.TestCase):
+	@classmethod
+	def setUpClass(cls):
+		pass
+
+	@classmethod
+	def tearDownClass(cls):
+		pass
+
+	# Returns a tuple of:
+	#   bool   : Two lists are equal
+	#   string : The string diff
+	# Args:
+	#   list
+	#   list   : (compare list one vs list two)
+	#   string : anything additional wanted to be printed before the string diff
+	#   bool   : the ordering of the items in the two lists must be the same
+	def report(self, oracle, test, firstText='', strictOrdering = False):
+		stringDiff = ""
+
+		equality = False
+		if not strictOrdering:
+			equality = (Counter(oracle) == Counter(test))
+		else:
+			equality = (oracle == test)
+
+		if equality:
+			return (True, '')
+		elif not strictOrdering:
+			try:
+					for elem in oracle:
+						test.remove(elem)
+						oracle.remove(elem)  # If it's not in the test, it won't get here!
+			except ValueError:
+					pass
+
+		differ = difflib.Differ(charjunk=difflib.IS_CHARACTER_JUNK)
+		skipped_lines = 0
+		for delta in differ.compare(oracle, test):
+			if delta[0] == ' ':
+					skipped_lines += 1
+					continue
+			if skipped_lines > 0:
+					stringDiff += "<---" + str(skipped_lines) + ' same lines--->\n'
+					skipped_lines = 0
+			delta = delta.replace('\n', '')
+			stringDiff += delta + '\n'
+
+		stringDiffList = stringDiff.split('\n')
+
+		if len(stringDiffList) > 10:
+			if not verbose:
+					stringDiff = '\n'.join(line if len(line) <= 100 else line[:100] + "...and " + str(len(line) - 100) + " more characters" for line in stringDiffList[:10])
+					stringDiff += '\n\n### And ' + str(len(stringDiffList)) + " more lines, use '-v' to show ###"
+		elif not verbose:
+			stringDiff = '\n'.join(line if len(line) <= 100 else line[:100] + "...and " + str(len(line) - 100) + " more characters" for line in stringDiffList)
+		stringDiff = '\n\n' + firstText + stringDiff
+		return (equality, stringDiff)
+
+	def run_rebase_test(self, testfile):
+		testname = None
+		with zipfile.ZipFile(os.path.join(api_suite_path, testfile), "r") as zf:
+			testname = zf.namelist()[0]
+			zf.extractall(path=api_suite_path)
+
+		pickle_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), testname + "_rebasing.pkl")
+		self.assertTrue(pickle_path, "Test pickle doesn't exist")
+		try:
+			# Python 2 does not have the encodings option
+			binary_oracle = pickle.load(open(pickle_path, "rb"), encoding='charmap')
+		except TypeError:
+			binary_oracle = pickle.load(open(pickle_path, "rb"))
+
+		test_builder = testcommon.BinaryViewTestBuilder(testname, imageBase=0xf00000)
+		self.assertTrue(test_builder.bv.start == 0xf00000)
+		for method in test_builder.methods():
+			test = getattr(test_builder, method)()
+			oracle = binary_oracle[method]
+			if test == oracle:
+					continue
+
+			result = getattr(test_builder, method).__doc__
+			result += ":\n"
+			report = self.report(oracle, test, result)
+			self.assertTrue(report[0], report[1])  # Test does not agree with oracle
+		os.unlink(os.path.join(api_suite_path, testname))
+
+	def test_rebasing__elf(self):
+		self.run_rebase_test('binaries/test_corpus/helloworld.zip')
+
+	def test_rebasing__macho(self):
+		self.run_rebase_test('binaries/test_corpus/duff.zip')
+
+	def test_rebasing__pe(self):
+		self.run_rebase_test('binaries/test_corpus/partial_register_dataflow.zip')
+
+	def test_rebasing__raw(self):
+		self.run_rebase_test('binaries/test_corpus/raw.zip')
+
+
 class MetaddataAPI(unittest.TestCase):
 	def test_metadata_basic_types(self):
 		# Core is tested thoroughly through the C++ unit tests here we focus on the python api side
diff --git a/suite/binaries b/suite/binaries
-Subproject 0f0801f1b75ac0e913ae9ca7a5447b18bf1c3a8
+Subproject 5440e9032ccf5e5d139ed7ba10a802c16ab2fbf
diff --git a/suite/generator.py b/suite/generator.py
index 443bb6bd..2fe7e8e0 100755
--- a/suite/generator.py
+++ b/suite/generator.py
@@ -290,6 +290,15 @@ def generate(test_store, outdir, exclude_binaries):
         binary_oracle.close()
         print("{0:.2f}".format(time.time() - binary_start_time))
 
+        # Generate oracle data for rebasing tests
+        name = oraclefile_rel[len(test_store):].replace(os.path.sep, "_").replace(".", "_")[1:]
+        if name in ["helloworld", "duff", "partial_register_dataflow", "raw"]:
+            test_data = testcommon.BinaryViewTestBuilder(oraclefile_rel, imageBase=0xf00000)
+            binary_oracle = OracleTestFile(oraclefile + "_rebasing")
+            for method in test_data.methods():
+                binary_oracle.add_entry(test_data, method)
+            binary_oracle.close()
+
         if not os.path.exists(oraclefile + ".zip"):
             with zipfile.ZipFile(oraclefile + ".zip", "w") as zf:
                 zf.write(oraclefile, os.path.relpath(oraclefile, start=os.path.dirname(__file__)))
diff --git a/suite/testcommon.py b/suite/testcommon.py
index 782f5e8d..6670d67c 100644
--- a/suite/testcommon.py
+++ b/suite/testcommon.py
@@ -98,9 +98,12 @@ class BinaryViewTestBuilder(Builder):
          - Function doc string used as 'on error' message
          - Should return: list of strings
     """
-    def __init__(self, filename):
+    def __init__(self, filename, imageBase=None):
         self.filename = os.path.join(os.path.dirname(__file__), filename)
-        self.bv = BinaryViewType.get_view_of_file(self.filename)
+        if imageBase is None:
+            self.bv = BinaryViewType.get_view_of_file(self.filename)
+        else:
+            self.bv = BinaryViewType.get_view_of_file_with_options(self.filename, options={'loader.imageBase' : imageBase})
         if self.bv is None:
             print("%s is not an executable format" % filename)
             return
diff --git a/unit_api.py b/unit_api.py
new file mode 100755
index 00000000..3c93e823
--- /dev/null
+++ b/unit_api.py
@@ -0,0 +1,27 @@
+#!/usr/bin/env python
+# This test file is maintained for API unit test distribution.
+import os
+import sys
+import unittest
+import pickle
+import zipfile
+import difflib
+from collections import Counter
+
+api_suite_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "api", "suite")
+sys.path.append(api_suite_path)
+import testcommon
+import api_test
+
+global verbose
+verbose = False
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1:
+        for i in range(1, len(sys.argv)):
+            if sys.argv[i] == '-v' or sys.argv[i] == '-V' or sys.argv[i] == '--verbose':
+                verbose = True
+
+    test_suite = unittest.defaultTestLoader.loadTestsFromModule(api_test)
+    runner = unittest.TextTestRunner(verbosity=2)
+    runner.run(test_suite)