Add ability to update custom flow graphs

8 files changed, 78 insertions, 9 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 1133625e..dd685d2b 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -2585,8 +2585,11 @@ namespace BinaryNinja
 		static void PrepareForLayoutCallback(void* ctxt);
 		static void PopulateNodesCallback(void* ctxt);
 		static void CompleteLayoutCallback(void* ctxt);
+		static BNFlowGraph* UpdateCallback(void* ctxt);
 
 	protected:
+		FlowGraph(BNFlowGraph* graph);
+
 		void FinishPrepareForLayout();
 		virtual void PrepareForLayout();
 		virtual void PopulateNodes();
@@ -2594,7 +2597,6 @@ namespace BinaryNinja
 
 	public:
 		FlowGraph();
-		FlowGraph(BNFlowGraph* graph);
 		~FlowGraph();
 
 		BNFlowGraph* GetGraphObject() const { return m_graph; }
@@ -2629,6 +2631,15 @@ namespace BinaryNinja
 		void SetMediumLevelILFunction(MediumLevelILFunction* func);
 
 		void Show(const std::string& title);
+
+		virtual Ref<FlowGraph> Update();
+	};
+
+	class CoreFlowGraph: public FlowGraph
+	{
+	public:
+		CoreFlowGraph(BNFlowGraph* graph);
+		virtual Ref<FlowGraph> Update() override;
 	};
 
 	struct LowLevelILLabel: public BNLowLevelILLabel
diff --git a/binaryninjacore.h b/binaryninjacore.h
index a32d8364..9359ca0e 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1778,6 +1778,7 @@ extern "C"
 		void (*prepareForLayout)(void* ctxt);
 		void (*populateNodes)(void* ctxt);
 		void (*completeLayout)(void* ctxt);
+		BNFlowGraph* (*update)(void* ctxt);
 	};
 
 	BINARYNINJACOREAPI char* BNAllocString(const char* contents);
@@ -2629,6 +2630,8 @@ extern "C"
 
 	BINARYNINJACOREAPI void BNFinishPrepareForLayout(BNFlowGraph* graph);
 
+	BINARYNINJACOREAPI BNFlowGraph* BNUpdateFlowGraph(BNFlowGraph* graph);
+
 	// Symbols
 	BINARYNINJACOREAPI BNSymbol* BNCreateSymbol(BNSymbolType type, const char* shortName, const char* fullName,
 	                                            const char* rawName, uint64_t addr);
diff --git a/flowgraph.cpp b/flowgraph.cpp
index 4bacee81..c779db0c 100644
--- a/flowgraph.cpp
+++ b/flowgraph.cpp
@@ -78,6 +78,16 @@ void FlowGraph::CompleteLayoutCallback(void* ctxt)
 }
 
 
+BNFlowGraph* FlowGraph::UpdateCallback(void* ctxt)
+{
+	FlowGraph* graph = (FlowGraph*)ctxt;
+	Ref<FlowGraph> result = graph->Update();
+	if (!result)
+		return nullptr;
+	return BNNewFlowGraphReference(result->GetGraphObject());
+}
+
+
 void FlowGraph::FinishPrepareForLayout()
 {
 	BNFinishPrepareForLayout(m_graph);
@@ -313,3 +323,23 @@ void FlowGraph::Show(const string& title)
 {
 	ShowGraphReport(title, this);
 }
+
+
+Ref<FlowGraph> FlowGraph::Update()
+{
+	return nullptr;
+}
+
+
+CoreFlowGraph::CoreFlowGraph(BNFlowGraph* graph): FlowGraph(graph)
+{
+}
+
+
+Ref<FlowGraph> CoreFlowGraph::Update()
+{
+	BNFlowGraph* graph = BNUpdateFlowGraph(GetGraphObject());
+	if (!graph)
+		return nullptr;
+	return new CoreFlowGraph(graph);
+}
diff --git a/function.cpp b/function.cpp
index 62a8e1b0..202f6d77 100644
--- a/function.cpp
+++ b/function.cpp
@@ -811,7 +811,7 @@ void Function::ApplyAutoDiscoveredType(Type* type)
 Ref<FlowGraph> Function::CreateFunctionGraph(BNFunctionGraphType type, DisassemblySettings* settings)
 {
 	BNFlowGraph* graph = BNCreateFunctionGraph(m_object, type, settings ? settings->GetObject() : nullptr);
-	return new FlowGraph(graph);
+	return new CoreFlowGraph(graph);
 }
 
 
@@ -1396,7 +1396,7 @@ Ref<FlowGraph> Function::GetUnresolvedStackAdjustmentGraph()
 	BNFlowGraph* graph = BNGetUnresolvedStackAdjustmentGraph(m_object);
 	if (!graph)
 		return nullptr;
-	return new FlowGraph(graph);
+	return new CoreFlowGraph(graph);
 }
 
 
diff --git a/interaction.cpp b/interaction.cpp
index 61c9dfb3..da3942bb 100644
--- a/interaction.cpp
+++ b/interaction.cpp
@@ -208,7 +208,7 @@ static void ShowGraphReportCallback(void* ctxt, BNBinaryView* view, const char*
 {
 	InteractionHandler* handler = (InteractionHandler*)ctxt;
 	handler->ShowGraphReport(view ? new BinaryView(BNNewViewReference(view)) : nullptr, title,
-		new FlowGraph(BNNewFlowGraphReference(graph)));
+		new CoreFlowGraph(BNNewFlowGraphReference(graph)));
 }
 
 
@@ -663,7 +663,7 @@ Ref<FlowGraph> ReportCollection::GetFlowGraph(size_t i) const
 	BNFlowGraph* graph = BNGetReportFlowGraph(m_object, i);
 	if (!graph)
 		return nullptr;
-	return new FlowGraph(graph);
+	return new CoreFlowGraph(graph);
 }
 
 
diff --git a/python/flowgraph.py b/python/flowgraph.py
index 25fdc430..e6c98597 100644
--- a/python/flowgraph.py
+++ b/python/flowgraph.py
@@ -268,6 +268,7 @@ class FlowGraph(object):
 			self._ext_cb.prepareForLayout = self._ext_cb.prepareForLayout.__class__(self._prepare_for_layout)
 			self._ext_cb.populateNodes = self._ext_cb.populateNodes.__class__(self._populate_nodes)
 			self._ext_cb.completeLayout = self._ext_cb.completeLayout.__class__(self._complete_layout)
+			self._ext_cb.update = self._ext_cb.update.__class__(self._update)
 			handle = core.BNCreateCustomFlowGraph(self._ext_cb)
 		self.handle = handle
 		self._on_complete = None
@@ -305,6 +306,16 @@ class FlowGraph(object):
 		except:
 			log.log_error(traceback.format_exc())
 
+	def _update(self, ctxt):
+		try:
+			graph = self.update()
+			if graph is None:
+				return None
+			return core.BNNewFlowGraphReference(graph.handle)
+		except:
+			log.log_error(traceback.format_exc())
+			return None
+
 	def finish_prepare_for_layout(self):
 		core.BNFinishPrepareForLayout(self.handle)
 
@@ -497,3 +508,17 @@ class FlowGraph(object):
 
 	def show(self, title):
 		interaction.show_graph_report(title, self)
+
+	def update(self):
+		return None
+
+
+class CoreFlowGraph(FlowGraph):
+	def __init__(self, handle):
+		super(CoreFlowGraph, self).__init__(handle)
+
+	def update(self):
+		graph = core.BNUpdateFlowGraph(self.handle)
+		if not graph:
+			return None
+		return CoreFlowGraph(graph)
diff --git a/python/function.py b/python/function.py
index dca9dda6..f64a9677 100644
--- a/python/function.py
+++ b/python/function.py
@@ -851,7 +851,7 @@ class Function(object):
 		graph = core.BNGetUnresolvedStackAdjustmentGraph(self.handle)
 		if not graph:
 			return None
-		return flowgraph.FlowGraph(graph)
+		return flowgraph.CoreFlowGraph(graph)
 
 	def __iter__(self):
 		count = ctypes.c_ulonglong()
@@ -1122,7 +1122,7 @@ class Function(object):
 			settings_obj = settings.handle
 		else:
 			settings_obj = None
-		return flowgraph.FlowGraph(core.BNCreateFunctionGraph(self.handle, graph_type, settings_obj))
+		return flowgraph.CoreFlowGraph(core.BNCreateFunctionGraph(self.handle, graph_type, settings_obj))
 
 	def apply_imported_types(self, sym):
 		core.BNApplyImportedTypes(self.handle, sym.handle)
diff --git a/python/interaction.py b/python/interaction.py
index 81aeb04f..b9aa0c2a 100644
--- a/python/interaction.py
+++ b/python/interaction.py
@@ -302,7 +302,7 @@ class InteractionHandler(object):
 				view = binaryview.BinaryView(handle = core.BNNewViewReference(view))
 			else:
 				view = None
-			self.show_graph_report(view, title, flowgraph.FlowGraph(core.BNNewFlowGraphReference(graph)))
+			self.show_graph_report(view, title, flowgraph.CoreFlowGraph(core.BNNewFlowGraphReference(graph)))
 		except:
 			log.log_error(traceback.format_exc())
 
@@ -567,7 +567,7 @@ class ReportCollection(object):
 			plaintext = core.BNGetReportPlainText(self.handle, i)
 			return HTMLReport(title, contents, plaintext, view)
 		elif report_type == ReportType.FlowGraphReportType:
-			graph = flowgraph.FlowGraph(core.BNGetReportFlowGraph(self.handle, i))
+			graph = flowgraph.CoreFlowGraph(core.BNGetReportFlowGraph(self.handle, i))
 			return FlowGraphReport(title, graph, view)
 		raise TypeError("invalid report type %s" % repr(report_type))