Add unit tests and update binaryview save

2 files changed, 82 insertions, 4 deletions

diff --git a/python/binaryview.py b/python/binaryview.py
index 146c2394..cefdabce 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -2318,28 +2318,30 @@ class BinaryView(object):
 		"""
 		return False
 
-	def create_database(self, filename, progress_func=None):
+	def create_database(self, filename, progress_func=None, clean=False):
 		"""
 		``create_database`` writes the current database (.bndb) file out to the specified file.
 
 		:param str filename: path and filename to write the bndb to, this string `should` have ".bndb" appended to it.
 		:param callback progress_func: optional function to be called with the current progress and total count.
+		:param bool clean: optional argument to determine if undo data is saved in the database.
 		:return: true on success, false on failure
 		:rtype: bool
 		"""
-		return self._file.create_database(filename, progress_func)
+		return self._file.create_database(filename, progress_func, clean)
 
-	def save_auto_snapshot(self, progress_func=None):
+	def save_auto_snapshot(self, progress_func=None, clean=False):
 		"""
 		``save_auto_snapshot`` saves the current database to the already created file.
 
 		.. note:: :py:meth:`create_database` should have been called prior to executing this method
 
 		:param callback progress_func: optional function to be called with the current progress and total count.
+		:param bool clean: optional argument to determine if undo data is saved in the database.
 		:return: True if it successfully saved the snapshot, False otherwise
 		:rtype: bool
 		"""
-		return self._file.save_auto_snapshot(progress_func)
+		return self._file.save_auto_snapshot(progress_func, clean)
 
 	def get_view_of_type(self, name):
 		"""
diff --git a/suite/testcommon.py b/suite/testcommon.py
index 06ef1b5b..628e129f 100644
--- a/suite/testcommon.py
+++ b/suite/testcommon.py
@@ -926,6 +926,82 @@ class VerifyBuilder(Builder):
         finally:
             self.delete_package("helloworld")
 
+    def test_verify_persistent_undo(self):
+        file_name = self.unpackage_file("helloworld")
+        try:
+            temp_name = next(tempfile._get_candidate_names()) + ".bndb"
+
+            bv = binja.BinaryViewType['ELF'].open(file_name)
+            bv.update_analysis_and_wait()
+
+            bv.begin_undo_actions()
+            bv.functions[0].set_comment(bv.functions[0].start, "Function start")
+            bv.functions[0].set_comment(bv.functions[0].start, "Function start!")
+            bv.commit_undo_actions()
+
+            bv.begin_undo_actions()
+            bv.add_function(bv.functions[0].start + 4)
+            bv.commit_undo_actions()
+
+            comments = self.get_comments(bv)
+            functions = self.get_functions(bv)
+
+            bv.create_database(temp_name)
+            bv.file.close()
+            del bv
+
+            bv = binja.FileMetadata(temp_name).open_existing_database(temp_name).get_view_of_type('ELF')
+            bv.update_analysis_and_wait()
+
+            bv.undo()
+            bv.undo()
+
+            bndb_functions = self.get_functions(bv)
+            bndb_comments = self.get_comments(bv)
+
+            bv.file.close()
+            del bv
+            os.unlink(temp_name)
+
+            return functions == bndb_functions and comments == bndb_comments
+        finally:
+            self.delete_package("helloworld")
+
+    def test_verify_clean_save(self):
+        file_name = self.unpackage_file("helloworld")
+        try:
+            temp_name = next(tempfile._get_candidate_names()) + ".bndb"
+
+            bv = binja.BinaryViewType['ELF'].open(file_name)
+            bv.update_analysis_and_wait()
+
+            bv.begin_undo_actions()
+            bv.functions[0].set_comment(bv.functions[0].start, "This is a secret comment")
+            bv.commit_undo_actions()
+
+            bv.begin_undo_actions()
+            bv.functions[0].set_comment(bv.functions[0].start, "Function start!")
+            bv.commit_undo_actions()
+
+            bv.create_database(temp_name, clean=True)
+            bv.file.close()
+            del bv
+
+            bv = binja.FileMetadata(temp_name).open_existing_database(temp_name).get_view_of_type('ELF')
+            bv.update_analysis_and_wait()
+
+            bv.undo()
+
+            comment = bv.functions[0].get_comment_at(bv.functions[0].start)
+
+            bv.file.close()
+            del bv
+            os.unlink(temp_name)
+
+            return comment == "Function start!"
+        finally:
+            self.delete_package("helloworld")
+
     def test_memory_leaks(self):
         """Detected memory leaks during analysis"""
         # This test will attempt to detect object leaks during headless analysis