Add analysis completion event and wait API

4 files changed, 102 insertions, 0 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 72240975..f7d260a1 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -571,6 +571,20 @@ namespace BinaryNinja
 		uint64_t addr;
 	};
 
+	class AnalysisCompletionEvent: public CoreRefCountObject<BNAnalysisCompletionEvent,
+		BNNewAnalysisCompletionEventReference, BNFreeAnalysisCompletionEvent>
+	{
+	protected:
+		std::function<void()> m_callback;
+		std::recursive_mutex m_mutex;
+
+		static void CompletionCallback(void* ctxt);
+
+	public:
+		AnalysisCompletionEvent(BinaryView* view, const std::function<void()>& callback);
+		void Cancel();
+	};
+
 	class BinaryView: public CoreRefCountObject<BNBinaryView, BNNewViewReference, BNFreeBinaryView>
 	{
 	protected:
@@ -740,6 +754,8 @@ namespace BinaryNinja
 
 		std::vector<BNStringReference> GetStrings();
 		std::vector<BNStringReference> GetStrings(uint64_t start, uint64_t len);
+
+		Ref<AnalysisCompletionEvent> AddAnalysisCompletionEvent(const std::function<void()>& callback);
 	};
 
 	class BinaryData: public BinaryView
diff --git a/binaryninjacore.h b/binaryninjacore.h
index cbbfcfa8..20f1b204 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -95,6 +95,7 @@ extern "C"
 	struct BNEnumeration;
 	struct BNCallingConvention;
 	struct BNPlatform;
+	struct BNAnalysisCompletionEvent;
 
 	enum BNLogLevel
 	{
@@ -1203,6 +1204,12 @@ extern "C"
 		uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI void BNFreeInstructionTextLines(BNInstructionTextLine* lines, size_t count);
 
+	BINARYNINJACOREAPI BNAnalysisCompletionEvent* BNAddAnalysisCompletionEvent(BNBinaryView* view, void* ctxt,
+		void (*callback)(void* ctxt));
+	BINARYNINJACOREAPI BNAnalysisCompletionEvent* BNNewAnalysisCompletionEventReference(BNAnalysisCompletionEvent* event);
+	BINARYNINJACOREAPI void BNFreeAnalysisCompletionEvent(BNAnalysisCompletionEvent* event);
+	BINARYNINJACOREAPI void BNCancelAnalysisCompletionEvent(BNAnalysisCompletionEvent* event);
+
 	// Function graph
 	BINARYNINJACOREAPI BNFunctionGraph* BNCreateFunctionGraph(BNFunction* func);
 	BINARYNINJACOREAPI BNFunctionGraph* BNNewFunctionGraphReference(BNFunctionGraph* graph);
diff --git a/binaryview.cpp b/binaryview.cpp
index 21a502ba..9a9461a4 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -180,6 +180,30 @@ Ref<Symbol> Symbol::ImportedFunctionFromImportAddressSymbol(Symbol* sym, uint64_
 }
 
 
+AnalysisCompletionEvent::AnalysisCompletionEvent(BinaryView* view, const std::function<void()>& callback):
+	m_callback(callback)
+{
+	m_object = BNAddAnalysisCompletionEvent(view->GetObject(), this, CompletionCallback);
+}
+
+
+void AnalysisCompletionEvent::CompletionCallback(void* ctxt)
+{
+	AnalysisCompletionEvent* event = (AnalysisCompletionEvent*)ctxt;
+
+	unique_lock<recursive_mutex> lock(event->m_mutex);
+	event->m_callback();
+	event->m_callback = []() {};
+}
+
+
+void AnalysisCompletionEvent::Cancel()
+{
+	unique_lock<recursive_mutex> lock(m_mutex);
+	m_callback = []() {};
+}
+
+
 BinaryView::BinaryView(const std::string& typeName, FileMetadata* file)
 {
 	BNCustomBinaryView view;
@@ -1064,6 +1088,12 @@ vector<BNStringReference> BinaryView::GetStrings(uint64_t start, uint64_t len)
 }
 
 
+Ref<AnalysisCompletionEvent> BinaryView::AddAnalysisCompletionEvent(const function<void()>& callback)
+{
+	return new AnalysisCompletionEvent(this, callback);
+}
+
+
 BinaryData::BinaryData(FileMetadata* file): BinaryView(BNCreateBinaryDataView(file->GetObject()))
 {
 }
diff --git a/python/__init__.py b/python/__init__.py
index 8024cae8..3d2123b3 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -627,6 +627,29 @@ class BinaryViewType(object):
 		except AttributeError:
 			raise AttributeError, "attribute '%s' is read only" % name
 
+class AnalysisCompletionEvent(object):
+	def __init__(self, view, callback):
+		self.view = view
+		self.callback = callback
+		self._cb = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self._notify)
+		self.handle = core.BNAddAnalysisCompletionEvent(self.view.handle, None, self._cb)
+
+	def __del__(self):
+		core.BNFreeAnalysisCompletionEvent(self.handle)
+
+	def _notify(self, ctxt):
+		try:
+			self.callback()
+		except:
+			log_error(traceback.format_exc())
+
+	def _empty_callback(self):
+		pass
+
+	def cancel(self):
+		self.callback = self._empty_callback
+		core.BNCancelAnalysisCompletionEvent(self.handle)
+
 class BinaryView(object):
 	name = None
 	"""Binary View name"""
@@ -1252,6 +1275,29 @@ class BinaryView(object):
 	def update_analysis(self):
 		core.BNUpdateAnalysis(self.handle)
 
+	def update_analysis_and_wait(self):
+		class WaitEvent:
+			def __init__(self):
+				self.cond = threading.Condition()
+				self.done = False
+
+			def complete(self):
+				self.cond.acquire()
+				self.done = True
+				self.cond.notify()
+				self.cond.release()
+
+			def wait(self):
+				self.cond.acquire()
+				while not self.done:
+					self.cond.wait()
+				self.cond.release()
+
+		wait = WaitEvent()
+		event = AnalysisCompletionEvent(self, lambda: wait.complete())
+		core.BNUpdateAnalysis(self.handle)
+		wait.wait()
+
 	def abort_analysis(self):
 		core.BNAbortAnalysis(self.handle)
 
@@ -1428,6 +1474,9 @@ class BinaryView(object):
 		core.BNFreeStringList(strings)
 		return result
 
+	def add_analysis_completion_event(self, callback):
+		return AnalysisCompletionEvent(self, callback)
+
 	def __setattr__(self, name, value):
 		try:
 			object.__setattr__(self,name,value)