Add support for custom flowgraph edges

9 files changed, 159 insertions, 105 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 05924070..ddfa6340 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -3230,6 +3230,7 @@ __attribute__ ((format (printf, 1, 2)))
 		Ref<FlowGraphNode> target;
 		std::vector<BNPoint> points;
 		bool backEdge;
+		BNEdgeStyle style;
 	};
 
 	class FlowGraphNode: public CoreRefCountObject<BNFlowGraphNode,
@@ -3255,7 +3256,7 @@ __attribute__ ((format (printf, 1, 2)))
 		void SetLines(const std::vector<DisassemblyTextLine>& lines);
 		const std::vector<FlowGraphEdge>& GetOutgoingEdges();
 		const std::vector<FlowGraphEdge>& GetIncomingEdges();
-		void AddOutgoingEdge(BNBranchType type, FlowGraphNode* target);
+		void AddOutgoingEdge(BNBranchType type, FlowGraphNode* target, BNEdgeStyle edgeStyle = BNEdgeStyle());
 
 		BNHighlightColor GetHighlight() const;
 		void SetHighlight(const BNHighlightColor& color);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 9cc2bb32..56f55360 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -222,7 +222,8 @@ extern "C"
 		SystemCall = 5,
 		IndirectBranch = 6,
 		ExceptionBranch = 7,
-		UnresolvedBranch = 127
+		UnresolvedBranch = 127,
+		UserDefinedBranch = 128
 	};
 
 	enum BNInstructionTextTokenType
@@ -1509,6 +1510,111 @@ extern "C"
 		float y;
 	};
 
+	enum BNThemeColor
+	{
+		// Hex dump colors
+		AddressColor,
+		ModifiedColor,
+		InsertedColor,
+		NotPresentColor,
+		SelectionColor,
+		OutlineColor,
+		BackgroundHighlightDarkColor,
+		BackgroundHighlightLightColor,
+		BoldBackgroundHighlightDarkColor,
+		BoldBackgroundHighlightLightColor,
+		AlphanumericHighlightColor,
+		PrintableHighlightColor,
+
+		// Graph colors
+		GraphBackgroundDarkColor,
+		GraphBackgroundLightColor,
+		GraphNodeDarkColor,
+		GraphNodeLightColor,
+		GraphNodeOutlineColor,
+		TrueBranchColor,
+		FalseBranchColor,
+		UnconditionalBranchColor,
+		AltTrueBranchColor,
+		AltFalseBranchColor,
+		AltUnconditionalBranchColor,
+
+		// Disassembly colors
+		RegisterColor,
+		NumberColor,
+		CodeSymbolColor,
+		DataSymbolColor,
+		StackVariableColor,
+		ImportColor,
+		InstructionHighlightColor,
+		TokenHighlightColor,
+		TokenSelectionColor,
+		AnnotationColor,
+		OpcodeColor,
+		LinearDisassemblyFunctionHeaderColor,
+		LinearDisassemblyBlockColor,
+		LinearDisassemblyNoteColor,
+		LinearDisassemblySeparatorColor,
+		StringColor,
+		TypeNameColor,
+		FieldNameColor,
+		KeywordColor,
+		UncertainColor,
+		NameSpaceColor,
+		NameSpaceSeparatorColor,
+		GotoLabelColor,
+		CommentColor,
+
+		// Script console colors
+		ScriptConsoleOutputColor,
+		ScriptConsoleErrorColor,
+		ScriptConsoleEchoColor,
+
+		// Highlighting colors
+		BlueStandardHighlightColor,
+		GreenStandardHighlightColor,
+		CyanStandardHighlightColor,
+		RedStandardHighlightColor,
+		MagentaStandardHighlightColor,
+		YellowStandardHighlightColor,
+		OrangeStandardHighlightColor,
+		WhiteStandardHighlightColor,
+		BlackStandardHighlightColor,
+
+		// MiniGraph
+		MiniGraphOverlayColor,
+
+		// FeatureMap
+		FeatureMapBaseColor,
+		FeatureMapNavLineColor,
+		FeatureMapNavHighlightColor,
+		FeatureMapDataVariableColor,
+		FeatureMapAsciiStringColor,
+		FeatureMapUnicodeStringColor,
+		FeatureMapFunctionColor,
+		FeatureMapImportColor,
+		FeatureMapExternColor,
+		FeatureMapLibraryColor
+	};
+
+	// The following edge styles map to QT's Qt::PenStyle enumeration
+	enum BNEdgePenStyle
+	{
+		NoPen = 0,          // no line at all.
+		SolidLine = 1,      // A plain line (default)
+		DashLine = 2,       // Dashes separated by a few pixels.
+		DotLine = 3,        // Dots separated by a few pixels.
+		DashDotLine = 4,    // Alternate dots and dashes.
+		DashDotDotLine = 5, // One dash, two dots, one dash, two dots.
+	};
+
+	struct BNEdgeStyle
+	{
+		BNEdgePenStyle style;
+		size_t width;
+		BNThemeColor color;
+	};
+
 	struct BNFlowGraphEdge
 	{
 		BNBranchType type;
@@ -1516,6 +1622,7 @@ extern "C"
 		BNPoint* points;
 		size_t pointCount;
 		bool backEdge;
+		BNEdgeStyle style;
 	};
 
 	enum BNHighlightColorStyle
@@ -3510,7 +3617,7 @@ __attribute__ ((format (printf, 1, 2)))
 	BINARYNINJACOREAPI BNFlowGraphEdge* BNGetFlowGraphNodeOutgoingEdges(BNFlowGraphNode* node, size_t* count);
 	BINARYNINJACOREAPI BNFlowGraphEdge* BNGetFlowGraphNodeIncomingEdges(BNFlowGraphNode* node, size_t* count);
 	BINARYNINJACOREAPI void BNFreeFlowGraphNodeEdgeList(BNFlowGraphEdge* edges, size_t count);
-	BINARYNINJACOREAPI void BNAddFlowGraphNodeOutgoingEdge(BNFlowGraphNode* node, BNBranchType type, BNFlowGraphNode* target);
+	BINARYNINJACOREAPI void BNAddFlowGraphNodeOutgoingEdge(BNFlowGraphNode* node, BNBranchType type, BNFlowGraphNode* target, BNEdgeStyle edgeStyle);
 
 	BINARYNINJACOREAPI BNHighlightColor BNGetFlowGraphNodeHighlight(BNFlowGraphNode* node);
 	BINARYNINJACOREAPI void BNSetFlowGraphNodeHighlight(BNFlowGraphNode* node, BNHighlightColor color);
diff --git a/flowgraphnode.cpp b/flowgraphnode.cpp
index 3a9359de..344dd8a3 100644
--- a/flowgraphnode.cpp
+++ b/flowgraphnode.cpp
@@ -155,6 +155,9 @@ const vector<FlowGraphEdge>& FlowGraphNode::GetOutgoingEdges()
 		edge.target = edges[i].target ? new FlowGraphNode(BNNewFlowGraphNodeReference(edges[i].target)) : nullptr;
 		edge.points.insert(edge.points.begin(), &edges[i].points[0], &edges[i].points[edges[i].pointCount]);
 		edge.backEdge = edges[i].backEdge;
+		edge.style.color = edges[i].style.color;
+		edge.style.width = edges[i].style.width;
+		edge.style.style = edges[i].style.style;
 		result.push_back(edge);
 	}
 
@@ -192,9 +195,9 @@ const vector<FlowGraphEdge>& FlowGraphNode::GetIncomingEdges()
 }
 
 
-void FlowGraphNode::AddOutgoingEdge(BNBranchType type, FlowGraphNode* target)
+void FlowGraphNode::AddOutgoingEdge(BNBranchType type, FlowGraphNode* target, BNEdgeStyle edgeStyle)
 {
-	BNAddFlowGraphNodeOutgoingEdge(m_object, type, target->GetObject());
+	BNAddFlowGraphNodeOutgoingEdge(m_object, type, target->GetObject(), edgeStyle);
 	m_cachedEdges.clear();
 	m_cachedEdgesValid = false;
 }
diff --git a/python/examples/triage/byte.py b/python/examples/triage/byte.py
index fe887fb4..6a3a9a10 100644
--- a/python/examples/triage/byte.py
+++ b/python/examples/triage/byte.py
@@ -5,8 +5,8 @@ from PySide2.QtWidgets import QAbstractScrollArea, QAbstractSlider
 from PySide2.QtGui import QPainter, QPalette, QFont
 from PySide2.QtCore import Qt, QTimer, QRect
 import binaryninjaui
-from binaryninjaui import View, ViewType, RenderContext, UIContext, ThemeColor, UIAction
-from binaryninja.enums import LinearDisassemblyLineType
+from binaryninjaui import View, ViewType, RenderContext, UIContext, UIAction
+from binaryninja.enums import LinearDisassemblyLineType, ThemeColor
 from binaryninja.binaryview import AddressRange
 
 
diff --git a/python/examples/triage/entropy.py b/python/examples/triage/entropy.py
index 57e9d2c4..3ce2de11 100644
--- a/python/examples/triage/entropy.py
+++ b/python/examples/triage/entropy.py
@@ -4,8 +4,8 @@ from PySide2.QtWidgets import QWidget
 from PySide2.QtGui import QImage, QColor, QPainter
 from PySide2.QtCore import Qt, QSize, QTimer
 import binaryninjaui
-from binaryninjaui import ViewFrame, ThemeColor, UIContext
-
+from binaryninjaui import ViewFrame, UIContext
+from binaryninja.enums import ThemeColor
 
 class EntropyThread(threading.Thread):
 	def __init__(self, data, image, block_size):
diff --git a/python/examples/triage/headers.py b/python/examples/triage/headers.py
index 33b9b3c2..0c87518d 100644
--- a/python/examples/triage/headers.py
+++ b/python/examples/triage/headers.py
@@ -1,7 +1,8 @@
 import time
 from binaryninja.binaryview import StructuredDataView
 import binaryninjaui
-from binaryninjaui import ThemeColor, ViewFrame, UIContext
+from binaryninjaui import ViewFrame, UIContext
+from binaryninja.enums import ThemeColor
 from PySide2.QtWidgets import QWidget, QLabel, QGridLayout
 from PySide2.QtGui import QPalette
 
diff --git a/python/examples/triage/sections.py b/python/examples/triage/sections.py
index 69ac8868..b18f5f28 100644
--- a/python/examples/triage/sections.py
+++ b/python/examples/triage/sections.py
@@ -3,7 +3,8 @@
 from PySide2.QtWidgets import QWidget, QLabel, QGridLayout, QHBoxLayout
 from binaryninja.enums import SectionSemantics
 import binaryninjaui
-from binaryninjaui import ThemeColor, ViewFrame, UIContext
+from binaryninjaui import ViewFrame, UIContext
+from binaryninja.enums import ThemeColor
 from . import headers
 
 
diff --git a/python/flowgraph.py b/python/flowgraph.py
index 645eab81..222095fb 100644
--- a/python/flowgraph.py
+++ b/python/flowgraph.py
@@ -24,7 +24,8 @@ import traceback
 
 # Binary Ninja components
 import binaryninja
-from binaryninja.enums import (BranchType, InstructionTextTokenType, HighlightStandardColor, FlowGraphOption)
+from binaryninja.enums import (BranchType, InstructionTextTokenType, HighlightStandardColor, FlowGraphOption,
+	EdgePenStyle, ThemeColor)
 from binaryninja import _binaryninjacore as core
 from binaryninja import function
 from binaryninja import binaryview
@@ -40,17 +41,35 @@ from binaryninja import range
 
 
 class FlowGraphEdge(object):
-	def __init__(self, branch_type, source, target, points, back_edge):
+	def __init__(self, branch_type, source, target, points, back_edge, style):
 		self.type = BranchType(branch_type)
 		self.source = source
 		self.target = target
 		self.points = points
 		self.back_edge = back_edge
+		self.style = style
 
 	def __repr__(self):
 		return "<%s: %s>" % (self.type.name, repr(self.target))
 
 
+class EdgeStyle(object):
+	def __init__(self, style=None, width=None, theme_color=None):
+		self.style = style if style is not None else EdgePenStyle.SolidLine
+		self.width = width if width is not None else 0
+		self.color = theme_color if theme_color is not None else ThemeColor.AddressColor
+
+	def _get_core_struct(self):
+		result = core.BNEdgeStyle()
+		result.style = self.style
+		result.width = self.width
+		result.color = self.color
+		return result
+
+	@classmethod
+	def from_core_struct(cls, edge_style):
+		return EdgeStyle(edge_style.style, edge_style.width, edge_style.color)
+
 class FlowGraphNode(object):
 	def __init__(self, graph = None, handle = None):
 		if handle is None:
@@ -228,7 +247,7 @@ class FlowGraphNode(object):
 			points = []
 			for j in range(0, edges[i].pointCount):
 				points.append((edges[i].points[j].x, edges[i].points[j].y))
-			result.append(FlowGraphEdge(branch_type, self, target, points, edges[i].backEdge))
+			result.append(FlowGraphEdge(branch_type, self, target, points, edges[i].backEdge, EdgeStyle(edges[i].style)))
 		core.BNFreeFlowGraphNodeEdgeList(edges, count.value)
 		return result
 
@@ -246,7 +265,7 @@ class FlowGraphNode(object):
 			points = []
 			for j in range(0, edges[i].pointCount):
 				points.append((edges[i].points[j].x, edges[i].points[j].y))
-			result.append(FlowGraphEdge(branch_type, self, target, points, edges[i].backEdge))
+			result.append(FlowGraphEdge(branch_type, self, target, points, edges[i].backEdge, EdgeStyle(edges[i].style)))
 		core.BNFreeFlowGraphNodeEdgeList(edges, count.value)
 		return result
 
@@ -271,16 +290,23 @@ class FlowGraphNode(object):
 			color = highlight.HighlightColor(color)
 		core.BNSetFlowGraphNodeHighlight(self.handle, color._get_core_struct())
 
-	def add_outgoing_edge(self, edge_type, target):
+	def add_outgoing_edge(self, edge_type, target, style=None):
 		"""
 		``add_outgoing_edge`` connects two flow graph nodes with an edge.
 
 		:param BranchType edge_type: Type of edge to add
 		:param FlowGraphNode target: Target node object
+		:param EdgeStyle style: (optional) Styling for graph edge Branch Type must be set to UserDefinedBranch
 		"""
 		if not target.is_valid_for_graph(self._graph):
 			raise ValueError("Target of edge has not been added to the owning graph")
-		core.BNAddFlowGraphNodeOutgoingEdge(self.handle, edge_type, target.handle)
+
+		if style is None:
+			style = EdgeStyle()
+		elif not isinstance(style, EdgeStyle):
+			raise AttributeError("style must be of type EdgeStyle")
+
+		core.BNAddFlowGraphNodeOutgoingEdge(self.handle, edge_type, target.handle, style._get_core_struct())
 
 	def is_valid_for_graph(self, graph):
 		return core.BNIsNodeValidForFlowGraph(graph.handle, self.handle)
@@ -338,7 +364,8 @@ class FlowGraph(object):
 		1
 		>>> graph.append(node_c)
 		2
-		>>> node_a.add_outgoing_edge(BranchType.UnconditionalBranch, node_b)
+		>>> edge = EdgeStyle(EdgePenStyle.DashDotDotLine, 2, ThemeColor.AddressColor)
+		>>> node_a.add_outgoing_edge(BranchType.UserDefinedBranch, node_b, edge)
 		>>> node_a.add_outgoing_edge(BranchType.UnconditionalBranch, node_c)
 		>>> show_graph_report("Custom Graph", graph)
 
diff --git a/ui/theme.h b/ui/theme.h
index 7010ab27..b4c1ead5 100644
--- a/ui/theme.h
+++ b/ui/theme.h
@@ -17,92 +17,6 @@ public:
 		const QWidget* widget) const override;
 };
 
-enum ThemeColor
-{
-	// Hex dump colors
-	AddressColor,
-	ModifiedColor,
-	InsertedColor,
-	NotPresentColor,
-	SelectionColor,
-	OutlineColor,
-	BackgroundHighlightDarkColor,
-	BackgroundHighlightLightColor,
-	BoldBackgroundHighlightDarkColor,
-	BoldBackgroundHighlightLightColor,
-	AlphanumericHighlightColor,
-	PrintableHighlightColor,
-
-	// Graph colors
-	GraphBackgroundDarkColor,
-	GraphBackgroundLightColor,
-	GraphNodeDarkColor,
-	GraphNodeLightColor,
-	GraphNodeOutlineColor,
-	TrueBranchColor,
-	FalseBranchColor,
-	UnconditionalBranchColor,
-	AltTrueBranchColor,
-	AltFalseBranchColor,
-	AltUnconditionalBranchColor,
-
-	// Disassembly colors
-	RegisterColor,
-	NumberColor,
-	CodeSymbolColor,
-	DataSymbolColor,
-	StackVariableColor,
-	ImportColor,
-	InstructionHighlightColor,
-	TokenHighlightColor,
-	TokenSelectionColor,
-	AnnotationColor,
-	OpcodeColor,
-	LinearDisassemblyFunctionHeaderColor,
-	LinearDisassemblyBlockColor,
-	LinearDisassemblyNoteColor,
-	LinearDisassemblySeparatorColor,
-	StringColor,
-	TypeNameColor,
-	FieldNameColor,
-	KeywordColor,
-	UncertainColor,
-	NameSpaceColor,
-	NameSpaceSeparatorColor,
-	GotoLabelColor,
-	CommentColor,
-
-	// Script console colors
-	ScriptConsoleOutputColor,
-	ScriptConsoleErrorColor,
-	ScriptConsoleEchoColor,
-
-	// Highlighting colors
-	BlueStandardHighlightColor,
-	GreenStandardHighlightColor,
-	CyanStandardHighlightColor,
-	RedStandardHighlightColor,
-	MagentaStandardHighlightColor,
-	YellowStandardHighlightColor,
-	OrangeStandardHighlightColor,
-	WhiteStandardHighlightColor,
-	BlackStandardHighlightColor,
-
-	// MiniGraph
-	MiniGraphOverlayColor,
-
-	// FeatureMap
-	FeatureMapBaseColor,
-	FeatureMapNavLineColor,
-	FeatureMapNavHighlightColor,
-	FeatureMapDataVariableColor,
-	FeatureMapAsciiStringColor,
-	FeatureMapUnicodeStringColor,
-	FeatureMapFunctionColor,
-	FeatureMapImportColor,
-	FeatureMapExternColor,
-	FeatureMapLibraryColor
-};
 
 void BINARYNINJAUIAPI initThemes();
 void BINARYNINJAUIAPI resetUserThemes();
@@ -114,7 +28,7 @@ void BINARYNINJAUIAPI setActiveTheme(const QString& name, bool saveToSettings =
 bool BINARYNINJAUIAPI isColorBlindMode();
 void BINARYNINJAUIAPI setColorBlindMode(bool active);
 
-QColor BINARYNINJAUIAPI getThemeColor(ThemeColor color);
+QColor BINARYNINJAUIAPI getThemeColor(BNThemeColor color);
 QColor BINARYNINJAUIAPI getTokenColor(QWidget* widget, BNInstructionTextTokenType token);
 
 QColor BINARYNINJAUIAPI avgColor(QColor a, QColor b);