Add flow graph options for additional UI features

4 files changed, 122 insertions, 1 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index c374d954..236bde65 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -2776,7 +2776,9 @@ namespace BinaryNinja
 		FlowGraph();
 
 		Ref<Function> GetFunction() const;
+		Ref<BinaryView> GetView() const;
 		void SetFunction(Function* func);
+		void SetView(BinaryView* view);
 
 		int GetHorizontalNodeMargin() const;
 		int GetVerticalNodeMargin() const;
@@ -2805,6 +2807,9 @@ namespace BinaryNinja
 		void Show(const std::string& title);
 
 		virtual Ref<FlowGraph> Update();
+
+		void SetOption(BNFlowGraphOption option, bool value = true);
+		bool IsOptionSet(BNFlowGraphOption option);
 	};
 
 	class CoreFlowGraph: public FlowGraph
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 2b831266..00163b92 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -702,6 +702,16 @@ extern "C"
 		PointerDisplayType
 	};
 
+	enum BNFlowGraphOption
+	{
+		FlowGraphUsesBlockHighlights,
+		FlowGraphUsesInstructionHighlights,
+		FlowGraphIncludesUserComments,
+		FlowGraphAllowsPatching,
+		FlowGraphAllowsInlineInstructionEditing,
+		FlowGraphShowsSecondaryRegisterHighlighting
+	};
+
 	struct BNLowLevelILInstruction
 	{
 		BNLowLevelILOperation operation;
@@ -2860,6 +2870,8 @@ extern "C"
 	BINARYNINJACOREAPI void BNFreeFlowGraph(BNFlowGraph* graph);
 	BINARYNINJACOREAPI BNFunction* BNGetFunctionForFlowGraph(BNFlowGraph* graph);
 	BINARYNINJACOREAPI void BNSetFunctionForFlowGraph(BNFlowGraph* graph, BNFunction* func);
+	BINARYNINJACOREAPI BNBinaryView* BNGetViewForFlowGraph(BNFlowGraph* graph);
+	BINARYNINJACOREAPI void BNSetViewForFlowGraph(BNFlowGraph* graph, BNBinaryView* view);
 
 	BINARYNINJACOREAPI int BNGetHorizontalFlowGraphNodeMargin(BNFlowGraph* graph);
 	BINARYNINJACOREAPI int BNGetVerticalFlowGraphNodeMargin(BNFlowGraph* graph);
@@ -2916,6 +2928,9 @@ extern "C"
 
 	BINARYNINJACOREAPI BNFlowGraph* BNUpdateFlowGraph(BNFlowGraph* graph);
 
+	BINARYNINJACOREAPI void BNSetFlowGraphOption(BNFlowGraph* graph, BNFlowGraphOption option, bool value);
+	BINARYNINJACOREAPI bool BNIsFlowGraphOptionSet(BNFlowGraph* graph, BNFlowGraphOption option);
+
 	// Symbols
 	BINARYNINJACOREAPI BNSymbol* BNCreateSymbol(BNSymbolType type, const char* shortName, const char* fullName,
 		const char* rawName, uint64_t addr, BNSymbolBinding binding, const BNNameSpace* nameSpace);
diff --git a/flowgraph.cpp b/flowgraph.cpp
index df1474cd..64bfe364 100644
--- a/flowgraph.cpp
+++ b/flowgraph.cpp
@@ -148,12 +148,27 @@ Ref<Function> FlowGraph::GetFunction() const
 }
 
 
+Ref<BinaryView> FlowGraph::GetView() const
+{
+	BNBinaryView* view = BNGetViewForFlowGraph(m_object);
+	if (!view)
+		return nullptr;
+	return new BinaryView(view);
+}
+
+
 void FlowGraph::SetFunction(Function* func)
 {
 	BNSetFunctionForFlowGraph(m_object, func ? func->GetObject() : nullptr);
 }
 
 
+void FlowGraph::SetView(BinaryView* view)
+{
+	BNSetViewForFlowGraph(m_object, view ? view->GetObject() : nullptr);
+}
+
+
 int FlowGraph::GetHorizontalNodeMargin() const
 {
 	return BNGetHorizontalFlowGraphNodeMargin(m_object);
@@ -344,6 +359,18 @@ Ref<FlowGraph> FlowGraph::Update()
 }
 
 
+void FlowGraph::SetOption(BNFlowGraphOption option, bool value)
+{
+	BNSetFlowGraphOption(m_object, option, value);
+}
+
+
+bool FlowGraph::IsOptionSet(BNFlowGraphOption option)
+{
+	return BNIsFlowGraphOptionSet(m_object, option);
+}
+
+
 CoreFlowGraph::CoreFlowGraph(BNFlowGraph* graph): FlowGraph(graph)
 {
 }
diff --git a/python/flowgraph.py b/python/flowgraph.py
index 4d2bd3ba..ce92e47b 100644
--- a/python/flowgraph.py
+++ b/python/flowgraph.py
@@ -24,7 +24,7 @@ import traceback
 
 # Binary Ninja components
 import binaryninja
-from binaryninja.enums import (BranchType, InstructionTextTokenType, HighlightStandardColor)
+from binaryninja.enums import (BranchType, InstructionTextTokenType, HighlightStandardColor, FlowGraphOption)
 from binaryninja import _binaryninjacore as core
 from binaryninja import function
 from binaryninja import binaryview
@@ -411,6 +411,20 @@ class FlowGraph(object):
 		core.BNSetFunctionForFlowGraph(self.handle, func)
 
 	@property
+	def view(self):
+		"""Binary view for a flow graph"""
+		view = core.BNGetViewForFlowGraph(self.handle)
+		if view is None:
+			return None
+		return binaryview.BinaryView(handle = view)
+
+	@view.setter
+	def view(self, view):
+		if view is not None:
+			view = view.handle
+		core.BNSetViewForFlowGraph(self.handle, view)
+
+	@property
 	def complete(self):
 		"""Whether flow graph layout is complete (read-only)"""
 		return core.BNIsFlowGraphLayoutComplete(self.handle)
@@ -503,6 +517,60 @@ class FlowGraph(object):
 		else:
 			raise TypeError("expected IL function for setting il_function property")
 
+	@property
+	def uses_block_highlights(self):
+		"""Set if flow graph uses the standard basic block highlighting settings"""
+		return self.is_option_set(FlowGraphOption.FlowGraphUsesBlockHighlights)
+
+	@uses_block_highlights.setter
+	def uses_block_highlights(self, value):
+		self.set_option(FlowGraphOption.FlowGraphUsesBlockHighlights, value)
+
+	@property
+	def uses_instruction_highlights(self):
+		"""Set if flow graph uses the standard instruction highlighting settings"""
+		return self.is_option_set(FlowGraphOption.FlowGraphUsesInstructionHighlights)
+
+	@uses_instruction_highlights.setter
+	def uses_instruction_highlights(self, value):
+		self.set_option(FlowGraphOption.FlowGraphUsesInstructionHighlights, value)
+
+	@property
+	def includes_user_comments(self):
+		"""Set if flow graph includes comments made by the user"""
+		return self.is_option_set(FlowGraphOption.FlowGraphIncludesUserComments)
+
+	@includes_user_comments.setter
+	def includes_user_comments(self, value):
+		self.set_option(FlowGraphOption.FlowGraphIncludesUserComments, value)
+
+	@property
+	def allows_patching(self):
+		"""Set if flow graph should allow modification of code from within the graph view"""
+		return self.is_option_set(FlowGraphOption.FlowGraphAllowsPatching)
+
+	@allows_patching.setter
+	def allows_patching(self, value):
+		self.set_option(FlowGraphOption.FlowGraphAllowsPatching, value)
+
+	@property
+	def allows_inline_instruction_editing(self):
+		"""Set if flow graph should allow inline instruction editing (assembly only)"""
+		return self.is_option_set(FlowGraphOption.FlowGraphAllowsInlineInstructionEditing)
+
+	@allows_inline_instruction_editing.setter
+	def allows_inline_instruction_editing(self, value):
+		self.set_option(FlowGraphOption.FlowGraphAllowsInlineInstructionEditing, value)
+
+	@property
+	def shows_secondary_reg_highlighting(self):
+		"""Set if flow graph should highlight associated registers in the UI"""
+		return self.is_option_set(FlowGraphOption.FlowGraphShowsSecondaryRegisterHighlighting)
+
+	@shows_secondary_reg_highlighting.setter
+	def shows_secondary_reg_highlighting(self, value):
+		self.set_option(FlowGraphOption.FlowGraphShowsSecondaryRegisterHighlighting, value)
+
 	def __setattr__(self, name, value):
 		try:
 			object.__setattr__(self, name, value)
@@ -604,6 +672,12 @@ class FlowGraph(object):
 		"""
 		return None
 
+	def set_option(self, option, value = True):
+		core.BNSetFlowGraphOption(self.handle, option, value)
+
+	def is_option_set(self, option):
+		return core.BNIsFlowGraphOptionSet(self.handle, option)
+
 
 class CoreFlowGraph(FlowGraph):
 	def __init__(self, handle):