From c5c93fc82b8929d04f62d241ca50228de60fa5f4 Mon Sep 17 00:00:00 2001
From: Rusty Wagner <rusty@vector35.com>
Date: Fri, 13 Jul 2018 18:43:34 -0400
Subject: Add ability to update custom flow graphs

---
 python/flowgraph.py   | 25 +++++++++++++++++++++++++
 python/function.py    |  4 ++--
 python/interaction.py |  4 ++--
 3 files changed, 29 insertions(+), 4 deletions(-)

(limited to 'python')

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))
 
-- 
cgit v1.3.1