Add cancelScript api for scriptingProvider

4 files changed, 48 insertions, 1 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index a79ad588..17069869 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -3997,6 +3997,7 @@ namespace BinaryNinja
 
 		static void DestroyInstanceCallback(void* ctxt);
 		static BNScriptingProviderExecuteResult ExecuteScriptInputCallback(void* ctxt, const char* input);
+		static void CancelScriptInputCallback(void* ctxt);
 		static void SetCurrentBinaryViewCallback(void* ctxt, BNBinaryView* view);
 		static void SetCurrentFunctionCallback(void* ctxt, BNFunction* func);
 		static void SetCurrentBasicBlockCallback(void* ctxt, BNBasicBlock* block);
@@ -4007,6 +4008,7 @@ namespace BinaryNinja
 
 	public:
 		virtual BNScriptingProviderExecuteResult ExecuteScriptInput(const std::string& input) = 0;
+		virtual void CancelScriptInput();
 		virtual void SetCurrentBinaryView(BinaryView* view);
 		virtual void SetCurrentFunction(Function* func);
 		virtual void SetCurrentBasicBlock(BasicBlock* block);
@@ -4029,6 +4031,7 @@ namespace BinaryNinja
 		virtual ~CoreScriptingInstance() {};
 
 		virtual BNScriptingProviderExecuteResult ExecuteScriptInput(const std::string& input) override;
+		virtual void CancelScriptInput() override;
 		virtual void SetCurrentBinaryView(BinaryView* view) override;
 		virtual void SetCurrentFunction(Function* func) override;
 		virtual void SetCurrentBasicBlock(BasicBlock* block) override;
diff --git a/binaryninjacore.h b/binaryninjacore.h
index d95b311c..ecc30c01 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1685,7 +1685,8 @@ extern "C"
 	{
 		InvalidScriptInput,
 		IncompleteScriptInput,
-		SuccessfulScriptExecution
+		SuccessfulScriptExecution,
+		ScriptExecutionCancelled
 	};
 
 
@@ -1694,6 +1695,7 @@ extern "C"
 		void* context;
 		void (*destroyInstance)(void* ctxt);
 		BNScriptingProviderExecuteResult (*executeScriptInput)(void* ctxt, const char* input);
+		void (*cancelScriptInput)(void* ctxt);
 		void (*setCurrentBinaryView)(void* ctxt, BNBinaryView* view);
 		void (*setCurrentFunction)(void* ctxt, BNFunction* func);
 		void (*setCurrentBasicBlock)(void* ctxt, BNBasicBlock* block);
@@ -3512,6 +3514,7 @@ extern "C"
 		BNScriptingInstance* instance);
 	BINARYNINJACOREAPI BNScriptingProviderExecuteResult BNExecuteScriptInput(BNScriptingInstance* instance,
 		const char* input);
+	BINARYNINJACOREAPI void BNCancelScriptInput(BNScriptingInstance* instance);
 	BINARYNINJACOREAPI void BNSetScriptingInstanceCurrentBinaryView(BNScriptingInstance* instance, BNBinaryView* view);
 	BINARYNINJACOREAPI void BNSetScriptingInstanceCurrentFunction(BNScriptingInstance* instance, BNFunction* func);
 	BINARYNINJACOREAPI void BNSetScriptingInstanceCurrentBasicBlock(BNScriptingInstance* instance, BNBasicBlock* block);
diff --git a/python/scriptingprovider.py b/python/scriptingprovider.py
index 21a68fcd..690bd6cc 100644
--- a/python/scriptingprovider.py
+++ b/python/scriptingprovider.py
@@ -107,6 +107,7 @@ class ScriptingInstance(object):
 			self._cb.context = 0
 			self._cb.destroyInstance = self._cb.destroyInstance.__class__(self._destroy_instance)
 			self._cb.executeScriptInput = self._cb.executeScriptInput.__class__(self._execute_script_input)
+			self._cb.cancelScriptInput = self._cb.cancelScriptInput.__class__(self._cancel_script_input)
 			self._cb.setCurrentBinaryView = self._cb.setCurrentBinaryView.__class__(self._set_current_binary_view)
 			self._cb.setCurrentFunction = self._cb.setCurrentFunction.__class__(self._set_current_function)
 			self._cb.setCurrentBasicBlock = self._cb.setCurrentBasicBlock.__class__(self._set_current_basic_block)
@@ -133,6 +134,13 @@ class ScriptingInstance(object):
 			log.log_error(traceback.format_exc())
 			return ScriptingProviderExecuteResult.InvalidScriptInput
 
+	def _cancel_script_input(self, ctxt):
+		try:
+			return self.perform_cancel_script_input()
+		except:
+			log.log_error(traceback.format_exc())
+			return ScriptingProviderExecuteResult.ScriptExecutionCancelled
+
 	def _set_current_binary_view(self, ctxt, view):
 		try:
 			if view:
@@ -189,6 +197,10 @@ class ScriptingInstance(object):
 		return ScriptingProviderExecuteResult.InvalidScriptInput
 
 	@abc.abstractmethod
+	def perform_cancel_script_input(self):
+		return ScriptingProviderExecuteResult.ScriptExecutionCancelled
+
+	@abc.abstractmethod
 	def perform_set_current_binary_view(self, view):
 		raise NotImplementedError
 
@@ -225,6 +237,9 @@ class ScriptingInstance(object):
 	def execute_script_input(self, text):
 		return core.BNExecuteScriptInput(self.handle, text)
 
+	def cancel_script_input(self, text):
+		return core.BNCancelScriptInput(self.handle)
+
 	def set_current_binary_view(self, view):
 		if view is not None:
 			view = view.handle
@@ -636,6 +651,14 @@ class PythonScriptingInstance(ScriptingInstance):
 		return ScriptingProviderExecuteResult.SuccessfulScriptExecution
 
 	@abc.abstractmethod
+	def perform_cancel_script_input(self):
+		for tid, tobj in threading._active.items():
+			if tobj is self.interpreter:
+				if ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(tid), ctypes.py_object(KeyboardInterrupt)) != 1:
+					ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(tid), None)
+				break
+
+	@abc.abstractmethod
 	def perform_set_current_binary_view(self, view):
 		self.interpreter.current_view = view
 
diff --git a/scriptingprovider.cpp b/scriptingprovider.cpp
index 2a233cfe..6483c040 100644
--- a/scriptingprovider.cpp
+++ b/scriptingprovider.cpp
@@ -55,6 +55,7 @@ ScriptingInstance::ScriptingInstance(ScriptingProvider* provider)
 	cb.context = this;
 	cb.destroyInstance = DestroyInstanceCallback;
 	cb.executeScriptInput = ExecuteScriptInputCallback;
+	cb.cancelScriptInput = CancelScriptInputCallback;
 	cb.setCurrentBinaryView = SetCurrentBinaryViewCallback;
 	cb.setCurrentFunction = SetCurrentFunctionCallback;
 	cb.setCurrentBasicBlock = SetCurrentBasicBlockCallback;
@@ -91,6 +92,13 @@ BNScriptingProviderExecuteResult ScriptingInstance::ExecuteScriptInputCallback(v
 }
 
 
+void ScriptingInstance::CancelScriptInputCallback(void* ctxt)
+{
+	ScriptingInstance* instance = (ScriptingInstance*)ctxt;
+	instance->CancelScriptInput();
+}
+
+
 void ScriptingInstance::SetCurrentBinaryViewCallback(void* ctxt, BNBinaryView* view)
 {
 	ScriptingInstance* instance = (ScriptingInstance*)ctxt;
@@ -131,6 +139,10 @@ void ScriptingInstance::DestroyInstance()
 }
 
 
+void ScriptingInstance::CancelScriptInput()
+{
+}
+
 void ScriptingInstance::SetCurrentBinaryView(BinaryView*)
 {
 }
@@ -203,6 +215,12 @@ BNScriptingProviderExecuteResult CoreScriptingInstance::ExecuteScriptInput(const
 }
 
 
+void CoreScriptingInstance::CancelScriptInput()
+{
+	BNCancelScriptInput(m_object);
+}
+
+
 void CoreScriptingInstance::SetCurrentBinaryView(BinaryView* view)
 {
 	BNSetScriptingInstanceCurrentBinaryView(m_object, view ? view->GetObject() : nullptr);