From dfe5d28f0aaee75ccad8c733e23879676ed77c37 Mon Sep 17 00:00:00 2001
From: Rusty Wagner <rusty@vector35.com>
Date: Wed, 11 Jul 2018 20:32:22 -0400
Subject: Add APIs for subclassing flow graphs

---
 python/flowgraph.py | 40 +++++++++++++++++++++++++++++++++++++---
 1 file changed, 37 insertions(+), 3 deletions(-)

(limited to 'python/flowgraph.py')

diff --git a/python/flowgraph.py b/python/flowgraph.py
index 1568bc77..25fdc430 100644
--- a/python/flowgraph.py
+++ b/python/flowgraph.py
@@ -24,13 +24,12 @@ import traceback
 
 # Binary Ninja components
 import _binaryninjacore as core
-from enums import (BranchType, InstructionTextTokenType, HighlightColorStyle, HighlightStandardColor)
+from enums import (BranchType, InstructionTextTokenType, HighlightStandardColor)
 import function
 import binaryview
 import lowlevelil
 import mediumlevelil
 import basicblock
-import architecture
 import log
 import interaction
 import highlight
@@ -264,7 +263,12 @@ class FlowGraphNode(object):
 class FlowGraph(object):
 	def __init__(self, handle = None):
 		if handle is None:
-			handle = core.BNCreateFlowGraph()
+			self._ext_cb = core.BNCustomFlowGraph()
+			self._ext_cb.context = 0
+			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)
+			handle = core.BNCreateCustomFlowGraph(self._ext_cb)
 		self.handle = handle
 		self._on_complete = None
 		self._cb = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self._complete)
@@ -283,6 +287,36 @@ class FlowGraph(object):
 			return True
 		return ctypes.addressof(self.handle.contents) != ctypes.addressof(value.handle.contents)
 
+	def _prepare_for_layout(self, ctxt):
+		try:
+			self.prepare_for_layout()
+		except:
+			log.log_error(traceback.format_exc())
+
+	def _populate_nodes(self, ctxt):
+		try:
+			self.populate_nodes()
+		except:
+			log.log_error(traceback.format_exc())
+
+	def _complete_layout(self, ctxt):
+		try:
+			self.complete_layout()
+		except:
+			log.log_error(traceback.format_exc())
+
+	def finish_prepare_for_layout(self):
+		core.BNFinishPrepareForLayout(self.handle)
+
+	def prepare_for_layout(self):
+		self.finish_prepare_for_layout()
+
+	def populate_nodes(self):
+		pass
+
+	def complete_layout(self):
+		pass
+
 	@property
 	def function(self):
 		"""Function for a flow graph"""
-- 
cgit v1.3.1