Undo transactions / context manager

7 files changed, 213 insertions, 46 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index d81c6320..5f4f6d16 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -2148,11 +2148,22 @@ namespace BinaryNinja {
 		bool CreateSnapshotedView(BinaryView* data, const std::string& viewName,
 								  const std::function<bool(size_t progress, size_t total)>& progressCallback);
 
+		/*! Run a function in a context in which any changes made to analysis will be added to an undo state.
+			If the function returns false or throws an exception, any changes made within will be reverted.
+
+			\param func Function to run in undo context
+			\return Return status of function
+			\throws std::exception If the called function throws an exception
+		 */
+		bool RunUndoableTransaction(std::function<bool()> func);
+
 		/*! Start recording actions taken so they can be undone at some point
 
+			\param anonymousAllowed Legacy interop: prevent empty calls to CommitUndoActions from affecting this
+			                        undo state. Specifically for RunUndoableTransaction.
 			\return Id of UndoEntry created, for passing to either CommitUndoActions or RevertUndoActions
 		*/
-		[[nodiscard]] std::string BeginUndoActions();
+		[[nodiscard]] std::string BeginUndoActions(bool anonymousAllowed = true);
 
 		/*!  Commit the actions taken since a call to BeginUndoActions.
 
@@ -3558,11 +3569,22 @@ namespace BinaryNinja {
 		bool SaveAutoSnapshot(const std::function<bool(size_t progress, size_t total)>& progressCallback,
 		    Ref<SaveSettings> settings = new SaveSettings());
 
+		/*! Run a function in a context in which any changes made to analysis will be added to an undo state.
+			If the function returns false or throws an exception, any changes made within will be reverted.
+
+			\param func Function to run in undo context
+			\return Return status of function
+			\throws std::exception If the called function throws an exception
+		 */
+		bool RunUndoableTransaction(std::function<bool()> func);
+
 		/*! Start recording actions taken so they can be undone at some point
 
+			\param anonymousAllowed Legacy interop: prevent empty calls to CommitUndoActions from affecting this
+			                        undo state. Specifically for RunUndoableTransaction.
 			\return Id of UndoEntry created, for passing to either CommitUndoActions or RevertUndoActions
 		*/
-		[[nodiscard]] std::string BeginUndoActions();
+		[[nodiscard]] std::string BeginUndoActions(bool anonymousAllowed = true);
 
 		/*!  Commit the actions taken since a call to BeginUndoActions.
 
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 815894f0..210a4322 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -3265,7 +3265,7 @@ extern "C"
 	BINARYNINJACOREAPI char* BNGetFilename(BNFileMetadata* file);
 	BINARYNINJACOREAPI void BNSetFilename(BNFileMetadata* file, const char* name);
 
-	BINARYNINJACOREAPI char* BNBeginUndoActions(BNFileMetadata* file);
+	BINARYNINJACOREAPI char* BNBeginUndoActions(BNFileMetadata* file, bool anonymousAllowed);
 	BINARYNINJACOREAPI void BNCommitUndoActions(BNFileMetadata* file, const char* id);
 	BINARYNINJACOREAPI void BNRevertUndoActions(BNFileMetadata* file, const char* id);
 
diff --git a/binaryview.cpp b/binaryview.cpp
index b5e80a68..dd54fbdd 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -1357,9 +1357,15 @@ bool BinaryView::SaveAutoSnapshot(
 }
 
 
-string BinaryView::BeginUndoActions()
+bool BinaryView::RunUndoableTransaction(std::function<bool()> func)
 {
-	return m_file->BeginUndoActions();
+	return m_file->RunUndoableTransaction(func);
+}
+
+
+string BinaryView::BeginUndoActions(bool anonymousAllowed)
+{
+	return m_file->BeginUndoActions(anonymousAllowed);
 }
 
 
diff --git a/filemetadata.cpp b/filemetadata.cpp
index 4f75129e..318da70d 100644
--- a/filemetadata.cpp
+++ b/filemetadata.cpp
@@ -270,9 +270,33 @@ bool FileMetadata::CreateSnapshotedView(BinaryView* data, const std::string& vie
 }
 
 
-std::string FileMetadata::BeginUndoActions()
+bool FileMetadata::RunUndoableTransaction(std::function<bool()> func)
 {
-	char* id = BNBeginUndoActions(m_object);
+	auto undo = BeginUndoActions(false);
+	try
+	{
+		bool result = func();
+		if (result)
+		{
+			CommitUndoActions(undo);
+		}
+		else
+		{
+			RevertUndoActions(undo);
+		}
+		return result;
+	}
+	catch (...)
+	{
+		RevertUndoActions(undo);
+		throw;
+	}
+}
+
+
+std::string FileMetadata::BeginUndoActions(bool anonymousAllowed)
+{
+	char* id = BNBeginUndoActions(m_object, anonymousAllowed);
 	std::string result = id;
 	BNFreeString(id);
 	return result;
diff --git a/python/binaryview.py b/python/binaryview.py
index 41fdf644..1e0dc9b4 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -3385,19 +3385,54 @@ class BinaryView:
 		"""
 		return self._file.get_view_of_type(name)
 
-	def begin_undo_actions(self) -> None:
+	def undoable_transaction(self) -> Generator:
 		"""
-		``begin_undo_actions`` start recording actions taken so the can be undone at some point.
+		``undoable_transaction`` gives you a context in which you can make changes to analysis,
+		and creates an Undo state containing those actions. If an exception is thrown, any
+		changes made to the analysis inside the transaction are reverted.
 
-		:rtype: None
+		:return: Transaction context manager, which will commit/revert actions depending on if an exception
+		         is thrown when it goes out of scope.
+		:rtype: Generator
+		:Example:
+
+			>>> bv.get_disassembly(0x100012f1)
+			'xor     eax, eax'
+			>>> # Actions inside the transaction will be committed to the undo state upon exit
+			>>> with bv.undoable_transaction():
+			>>>     bv.convert_to_nop(0x100012f1)
+			True
+			>>> bv.get_disassembly(0x100012f1)
+			'nop'
+			>>> bv.undo()
+			>>> bv.get_disassembly(0x100012f1)
+			'xor     eax, eax'
+			>>> # A thrown exception inside the transaction will undo all changes made inside it
+			>>> with bv.undoable_transaction():
+			>>>     bv.convert_to_nop(0x100012f1)  # Reverted on thrown exception
+			>>>     raise RuntimeError("oh no")
+			RuntimeError: oh no
+			>>> bv.get_disassembly(0x100012f1)
+			'xor     eax, eax'
+		"""
+		return self._file.undoable_transaction()
+
+	def begin_undo_actions(self, anonymous_allowed: bool = True) -> str:
+		"""
+		``begin_undo_actions`` starts recording actions taken so they can be undone at some point.
+
+		:param bool anonymous_allowed: Legacy interop: prevent empty calls to :py:func:`commit_undo_actions`` from
+		                               affecting this undo state. Specifically for :py:func:`undoable_transaction``
+		:return: Id of undo state, for passing to :py:func:`commit_undo_actions`` or :py:func:`revert_undo_actions`.
+		:rtype: str
 		:Example:
 
 			>>> bv.get_disassembly(0x100012f1)
 			'xor     eax, eax'
-			>>> bv.begin_undo_actions()
+			>>> state = bv.begin_undo_actions()
 			>>> bv.convert_to_nop(0x100012f1)
 			True
-			>>> bv.commit_undo_actions()
+			>>> bv.commit_undo_actions(state)
 			>>> bv.get_disassembly(0x100012f1)
 			'nop'
 			>>> bv.undo()
@@ -3405,21 +3440,24 @@ class BinaryView:
 			'xor     eax, eax'
 			>>>
 		"""
-		self._file.begin_undo_actions()
+		return self._file.begin_undo_actions(anonymous_allowed)
 
-	def commit_undo_actions(self) -> None:
+	def commit_undo_actions(self, id: Optional[str] = None) -> None:
 		"""
-		``commit_undo_actions`` commit the actions taken since the last commit to the undo database.
+		``commit_undo_actions`` commits the actions taken since a call to :py:func:`begin_undo_actions`
+		Pass as `id` the value returned by :py:func:`begin_undo_actions`. Empty values of
+		`id` will commit all changes since the last call to :py:func:`begin_undo_actions`.
 
+		:param Optional[str] id: id of undo state, from :py:func:`begin_undo_actions`
 		:rtype: None
 		:Example:
 
 			>>> bv.get_disassembly(0x100012f1)
 			'xor     eax, eax'
-			>>> bv.begin_undo_actions()
+			>>> state = bv.begin_undo_actions()
 			>>> bv.convert_to_nop(0x100012f1)
 			True
-			>>> bv.commit_undo_actions()
+			>>> bv.commit_undo_actions(state)
 			>>> bv.get_disassembly(0x100012f1)
 			'nop'
 			>>> bv.undo()
@@ -3427,21 +3465,42 @@ class BinaryView:
 			'xor     eax, eax'
 			>>>
 		"""
-		self._file.commit_undo_actions()
+		self._file.commit_undo_actions(id)
 
-	def undo(self) -> None:
+	def revert_undo_actions(self, id: Optional[str] = None) -> None:
 		"""
-		``undo`` undo the last committed action in the undo database.
+		``revert_undo_actions`` reverts the actions taken since a call to :py:func:`begin_undo_actions`
+		Pass as `id` the value returned by :py:func:`begin_undo_actions`. Empty values of
+		`id` will revert all changes since the last call to :py:func:`begin_undo_actions`.
 
+		:param Optional[str] id: id of undo state, from :py:func:`begin_undo_actions`
 		:rtype: None
 		:Example:
 
 			>>> bv.get_disassembly(0x100012f1)
 			'xor     eax, eax'
-			>>> bv.begin_undo_actions()
+			>>> state = bv.begin_undo_actions()
 			>>> bv.convert_to_nop(0x100012f1)
 			True
-			>>> bv.commit_undo_actions()
+			>>> bv.revert_undo_actions(state)
+			>>> bv.get_disassembly(0x100012f1)
+			'xor     eax, eax'
+			>>>
+		"""
+		self._file.revert_undo_actions(id)
+
+	def undo(self) -> None:
+		"""
+		``undo`` undo the last committed transaction in the undo database.
+
+		:rtype: None
+		:Example:
+
+			>>> bv.get_disassembly(0x100012f1)
+			'xor     eax, eax'
+			>>> with bv.undoable_transaction():
+			>>>     bv.convert_to_nop(0x100012f1)
+			True
 			>>> bv.get_disassembly(0x100012f1)
 			'nop'
 			>>> bv.undo()
@@ -3456,17 +3515,16 @@ class BinaryView:
 
 	def redo(self) -> None:
 		"""
-		``redo`` redo the last committed action in the undo database.
+		``redo`` redo the last committed transaction in the undo database.
 
 		:rtype: None
 		:Example:
 
 			>>> bv.get_disassembly(0x100012f1)
 			'xor     eax, eax'
-			>>> bv.begin_undo_actions()
-			>>> bv.convert_to_nop(0x100012f1)
+			>>> with bv.undoable_transaction():
+			>>>     bv.convert_to_nop(0x100012f1)
 			True
-			>>> bv.commit_undo_actions()
 			>>> bv.get_disassembly(0x100012f1)
 			'nop'
 			>>> bv.undo()
diff --git a/python/filemetadata.py b/python/filemetadata.py
index bdb46a00..f0d9391f 100644
--- a/python/filemetadata.py
+++ b/python/filemetadata.py
@@ -17,10 +17,10 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 # IN THE SOFTWARE.
-
+import contextlib
 import traceback
 import ctypes
-from typing import Any, Callable, Optional, List
+from typing import Any, Callable, Optional, List, Generator
 
 # Binary Ninja Components
 import binaryninja
@@ -304,20 +304,61 @@ class FileMetadata:
 		"""
 		core.BNCloseFile(self.handle)
 
-	def begin_undo_actions(self) -> str:
+	@contextlib.contextmanager
+	def undoable_transaction(self) -> Generator:
+		"""
+		``undoable_transaction`` gives you a context in which you can make changes to analysis,
+		and creates an Undo state containing those actions. If an exception is thrown, any
+		changes made to the analysis inside the transaction are reverted.
+
+		:return: Transaction context manager, which will commit/revert actions depending on if an exception
+		         is thrown when it goes out of scope.
+		:rtype: Generator
+		:Example:
+
+			>>> bv.get_disassembly(0x100012f1)
+			'xor     eax, eax'
+			>>> # Actions inside the transaction will be committed to the undo state upon exit
+			>>> with bv.undoable_transaction():
+			>>>     bv.convert_to_nop(0x100012f1)
+			True
+			>>> bv.get_disassembly(0x100012f1)
+			'nop'
+			>>> bv.undo()
+			>>> bv.get_disassembly(0x100012f1)
+			'xor     eax, eax'
+			>>> # A thrown exception inside the transaction will undo all changes made inside it
+			>>> with bv.undoable_transaction():
+			>>>     bv.convert_to_nop(0x100012f1)  # Reverted on thrown exception
+			>>>     raise RuntimeError("oh no")
+			RuntimeError: oh no
+			>>> bv.get_disassembly(0x100012f1)
+			'xor     eax, eax'
+		"""
+		state = self.begin_undo_actions(False)
+		try:
+			yield state
+			self.commit_undo_actions(state)
+		except:
+			self.revert_undo_actions(state)
+			raise
+
+	def begin_undo_actions(self, anonymous_allowed: bool = True) -> str:
 		"""
 		``begin_undo_actions`` starts recording actions taken so they can be undone at some point.
 
+		:param bool anonymous_allowed: Legacy interop: prevent empty calls to :py:func:`commit_undo_actions`` from
+		                               affecting this undo state. Specifically for :py:func:`undoable_transaction``
 		:return: Id of undo state, for passing to :py:func:`commit_undo_actions`` or :py:func:`revert_undo_actions`.
 		:rtype: str
 		:Example:
 
 			>>> bv.get_disassembly(0x100012f1)
 			'xor     eax, eax'
-			>>> undo = bv.begin_undo_actions()
+			>>> state = bv.begin_undo_actions()
 			>>> bv.convert_to_nop(0x100012f1)
 			True
-			>>> bv.commit_undo_actions(undo)
+			>>> bv.commit_undo_actions(state)
 			>>> bv.get_disassembly(0x100012f1)
 			'nop'
 			>>> bv.undo()
@@ -325,7 +366,7 @@ class FileMetadata:
 			'xor     eax, eax'
 			>>>
 		"""
-		id = core.BNBeginUndoActions(self.handle)
+		id = core.BNBeginUndoActions(self.handle, anonymous_allowed)
 		self._previous_undos.append(id)
 		return id
 
@@ -341,10 +382,10 @@ class FileMetadata:
 
 			>>> bv.get_disassembly(0x100012f1)
 			'xor     eax, eax'
-			>>> undo = bv.begin_undo_actions()
+			>>> state = bv.begin_undo_actions()
 			>>> bv.convert_to_nop(0x100012f1)
 			True
-			>>> bv.commit_undo_actions(undo)
+			>>> bv.commit_undo_actions(state)
 			>>> bv.get_disassembly(0x100012f1)
 			'nop'
 			>>> bv.undo()
@@ -369,10 +410,10 @@ class FileMetadata:
 
 			>>> bv.get_disassembly(0x100012f1)
 			'xor     eax, eax'
-			>>> undo = bv.begin_undo_actions()
+			>>> state = bv.begin_undo_actions()
 			>>> bv.convert_to_nop(0x100012f1)
 			True
-			>>> bv.revert_undo_actions(undo)
+			>>> bv.revert_undo_actions(state)
 			>>> bv.get_disassembly(0x100012f1)
 			'xor     eax, eax'
 			>>>
@@ -384,17 +425,16 @@ class FileMetadata:
 
 	def undo(self) -> None:
 		"""
-		``undo`` undo the last committed action in the undo database.
+		``undo`` undo the last committed transaction in the undo database.
 
 		:rtype: None
 		:Example:
 
 			>>> bv.get_disassembly(0x100012f1)
 			'xor     eax, eax'
-			>>> bv.begin_undo_actions()
-			>>> bv.convert_to_nop(0x100012f1)
+			>>> with bv.undoable_transaction():
+			>>>     bv.convert_to_nop(0x100012f1)
 			True
-			>>> bv.commit_undo_actions()
 			>>> bv.get_disassembly(0x100012f1)
 			'nop'
 			>>> bv.undo()
@@ -409,17 +449,16 @@ class FileMetadata:
 
 	def redo(self) -> None:
 		"""
-		``redo`` redo the last committed action in the undo database.
+		``redo`` redo the last committed transaction in the undo database.
 
 		:rtype: None
 		:Example:
 
 			>>> bv.get_disassembly(0x100012f1)
 			'xor     eax, eax'
-			>>> bv.begin_undo_actions()
-			>>> bv.convert_to_nop(0x100012f1)
+			>>> with bv.undoable_transaction():
+			>>>     bv.convert_to_nop(0x100012f1)
 			True
-			>>> bv.commit_undo_actions()
 			>>> bv.get_disassembly(0x100012f1)
 			'nop'
 			>>> bv.undo()
diff --git a/rust/src/filemetadata.rs b/rust/src/filemetadata.rs
index b0e1c73f..1a558f3d 100644
--- a/rust/src/filemetadata.rs
+++ b/rust/src/filemetadata.rs
@@ -125,8 +125,26 @@ impl FileMetadata {
         unsafe { BNIsBackedByDatabase(self.handle, view_type.as_ref().as_ptr() as *const _) }
     }
 
-    pub fn begin_undo_actions(&self) -> BnString {
-        unsafe { BnString::from_raw(BNBeginUndoActions(self.handle)) }
+    pub fn run_undoable_transaction<F: FnOnce() -> Result<T, E>, T, E>(
+        &self,
+        func: F,
+    ) -> Result<T, E> {
+        let undo = self.begin_undo_actions(false);
+        let result = func();
+        match result {
+            Ok(t) => {
+                self.commit_undo_actions(undo);
+                Ok(t)
+            }
+            Err(e) => {
+                self.revert_undo_actions(undo);
+                Err(e)
+            }
+        }
+    }
+
+    pub fn begin_undo_actions(&self, anonymous_allowed: bool) -> BnString {
+        unsafe { BnString::from_raw(BNBeginUndoActions(self.handle, anonymous_allowed)) }
     }
 
     pub fn commit_undo_actions<S: BnStrCompatible>(&self, id: S) {