Add APIs to display busy indicator for long running plugin tasks

4 files changed, 208 insertions, 0 deletions

diff --git a/backgroundtask.cpp b/backgroundtask.cpp
new file mode 100644
index 00000000..aebf980c
--- /dev/null
+++ b/backgroundtask.cpp
@@ -0,0 +1,75 @@
+#include "binaryninjaapi.h"
+
+using namespace std;
+using namespace BinaryNinja;
+
+
+BackgroundTask::BackgroundTask(BNBackgroundTask* task)
+{
+	m_object = task;
+}
+
+
+BackgroundTask::BackgroundTask(const string& initialText, bool canCancel)
+{
+	m_object = BNBeginBackgroundTask(initialText.c_str(), canCancel);
+}
+
+
+bool BackgroundTask::CanCancel() const
+{
+	return BNCanCancelBackgroundTask(m_object);
+}
+
+
+bool BackgroundTask::IsCancelled() const
+{
+	return BNIsBackgroundTaskCancelled(m_object);
+}
+
+
+bool BackgroundTask::IsFinished() const
+{
+	return BNIsBackgroundTaskFinished(m_object);
+}
+
+
+string BackgroundTask::GetProgressText() const
+{
+	char* text = BNGetBackgroundTaskProgressText(m_object);
+	string result = text;
+	BNFreeString(text);
+	return result;
+}
+
+
+void BackgroundTask::Cancel()
+{
+	BNCancelBackgroundTask(m_object);
+}
+
+
+void BackgroundTask::Finish()
+{
+	BNFinishBackgroundTask(m_object);
+}
+
+
+void BackgroundTask::SetProgressText(const string& text)
+{
+	BNSetBackgroundTaskProgressText(m_object, text.c_str());
+}
+
+
+vector<Ref<BackgroundTask>> BackgroundTask::GetRunningTasks()
+{
+	size_t count;
+	BNBackgroundTask** tasks = BNGetRunningBackgroundTasks(&count);
+
+	vector<Ref<BackgroundTask>> result;
+	for (size_t i = 0; i < count; i++)
+		result.push_back(new BackgroundTask(BNNewBackgroundTaskReference(tasks[i])));
+
+	BNFreeBackgroundTaskList(tasks, count);
+	return result;
+}
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 30285380..3400cca7 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -2260,4 +2260,23 @@ namespace BinaryNinja
 	public:
 		virtual void AddMainThreadAction(MainThreadAction* action) = 0;
 	};
+
+	class BackgroundTask: public CoreRefCountObject<BNBackgroundTask,
+		BNNewBackgroundTaskReference, BNFreeBackgroundTask>
+	{
+	public:
+		BackgroundTask(BNBackgroundTask* task);
+		BackgroundTask(const std::string& initialText, bool canCancel);
+
+		bool CanCancel() const;
+		bool IsCancelled() const;
+		bool IsFinished() const;
+		std::string GetProgressText() const;
+
+		void Cancel();
+		void Finish();
+		void SetProgressText(const std::string& text);
+
+		static std::vector<Ref<BackgroundTask>> GetRunningTasks();
+	};
 }
diff --git a/binaryninjacore.h b/binaryninjacore.h
index de7b77fd..89dceba1 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -102,6 +102,7 @@ extern "C"
 	struct BNScriptingProvider;
 	struct BNScriptingInstance;
 	struct BNMainThreadAction;
+	struct BNBackgroundTask;
 
 	//! Console log levels
 	enum BNLogLevel
@@ -2014,6 +2015,21 @@ extern "C"
 	BINARYNINJACOREAPI size_t BNGetWorkerThreadCount(void);
 	BINARYNINJACOREAPI void BNSetWorkerThreadCount(size_t count);
 
+	// Background task progress reporting
+	BINARYNINJACOREAPI BNBackgroundTask* BNBeginBackgroundTask(const char* initialText, bool canCancel);
+	BINARYNINJACOREAPI void BNFinishBackgroundTask(BNBackgroundTask* task);
+	BINARYNINJACOREAPI void BNSetBackgroundTaskProgressText(BNBackgroundTask* task, const char* text);
+	BINARYNINJACOREAPI bool BNIsBackgroundTaskCancelled(BNBackgroundTask* task);
+
+	BINARYNINJACOREAPI BNBackgroundTask** BNGetRunningBackgroundTasks(size_t* count);
+	BINARYNINJACOREAPI BNBackgroundTask* BNNewBackgroundTaskReference(BNBackgroundTask* task);
+	BINARYNINJACOREAPI void BNFreeBackgroundTask(BNBackgroundTask* task);
+	BINARYNINJACOREAPI void BNFreeBackgroundTaskList(BNBackgroundTask** tasks, size_t count);
+	BINARYNINJACOREAPI char* BNGetBackgroundTaskProgressText(BNBackgroundTask* task);
+	BINARYNINJACOREAPI bool BNCanCancelBackgroundTask(BNBackgroundTask* task);
+	BINARYNINJACOREAPI void BNCancelBackgroundTask(BNBackgroundTask* task);
+	BINARYNINJACOREAPI bool BNIsBackgroundTaskFinished(BNBackgroundTask* task);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/python/__init__.py b/python/__init__.py
index b29c9f80..6abbdfa6 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -9961,6 +9961,104 @@ class MainThreadActionHandler(object):
 	def add_action(self, action):
 		pass
 
+class _BackgroundTaskMetaclass(type):
+	@property
+	def list(self):
+		"""List all running background tasks (read-only)"""
+		count = ctypes.c_ulonglong()
+		tasks = core.BNGetRunningBackgroundTasks(count)
+		result = []
+		for i in xrange(0, count.value):
+			result.append(BackgroundTask(core.BNNewBackgroundTaskReference(tasks[i])))
+		core.BNFreeBackgroundTaskList(tasks)
+		return result
+
+	def __iter__(self):
+		_init_plugins()
+		count = ctypes.c_ulonglong()
+		tasks = core.BNGetRunningBackgroundTasks(count)
+		try:
+			for i in xrange(0, count.value):
+				yield BackgroundTask(core.BNNewBackgroundTaskReference(tasks[i]))
+		finally:
+			core.BNFreeBackgroundTaskList(tasks)
+
+class BackgroundTask(object):
+	__metaclass__ = _BackgroundTaskMetaclass
+
+	def __init__(self, initial_progress_text = "", can_cancel = False, handle = None):
+		if handle is None:
+			self.handle = core.BNBeginBackgroundTask(initial_progress_text, can_cancel)
+		else:
+			self.handle = handle
+
+	def __del__(self):
+		core.BNFreeBackgroundTask(self.handle)
+
+	@property
+	def progress(self):
+		"""Text description of the progress of the background task (displayed in status bar of the UI)"""
+		return core.BNGetBackgroundTaskProgressText(self.handle)
+
+	@progress.setter
+	def progress(self, value):
+		core.BNSetBackgroundTaskProgressText(self.handle, str(value))
+
+	@property
+	def can_cancel(self):
+		"""Whether the task can be cancelled (read-only)"""
+		return core.BNCanCancelBackgroundTask(self.handle)
+
+	@property
+	def finished(self):
+		"""Whether the task has finished"""
+		return core.BNIsBackgroundTaskFinished(self.handle)
+
+	@finished.setter
+	def finished(self, value):
+		if value:
+			self.finish()
+
+	def finish(self):
+		core.BNFinishBackgroundTask(self.handle)
+
+	@property
+	def cancelled(self):
+		"""Whether the task has been cancelled"""
+		return core.BNIsBackgroundTaskCancelled(self.handle)
+
+	@cancelled.setter
+	def cancelled(self, value):
+		if value:
+			self.cancel()
+
+	def cancel(self):
+		core.BNCancelBackgroundTask(self.handle)
+
+class BackgroundTaskThread(BackgroundTask):
+	def __init__(self, initial_progress_text = "", can_cancel = False):
+		class _Thread(threading.Thread):
+			def __init__(self, task):
+				threading.Thread.__init__(self)
+				self.task = task
+
+			def run(self):
+				self.task.run()
+				self.task.finish()
+				self.task = None
+
+		BackgroundTask.__init__(self, initial_progress_text, can_cancel)
+		self.thread = _Thread(self)
+
+	def run(self):
+		pass
+
+	def start(self):
+		self.thread.start()
+
+	def join(self):
+		self.thread.join()
+
 def LLIL_TEMP(n):
 	return n | 0x80000000