Early HLIL testing

1 files changed, 20 insertions, 0 deletions

diff --git a/python/flowgraph.py b/python/flowgraph.py
index 32303c93..6c3f41f7 100644
--- a/python/flowgraph.py
+++ b/python/flowgraph.py
@@ -30,6 +30,7 @@ from binaryninja import function
 from binaryninja import binaryview
 from binaryninja import lowlevelil
 from binaryninja import mediumlevelil
+from binaryninja import highlevelil
 from binaryninja import basicblock
 from binaryninja import log
 from binaryninja import highlight
@@ -532,6 +533,10 @@ class FlowGraph(object):
 		return core.BNIsMediumLevelILFlowGraph(self.handle)
 
 	@property
+	def is_high_level_il(self):
+		return core.BNIsHighLevelILFlowGraph(self.handle)
+
+	@property
 	def il_function(self):
 		if self.is_low_level_il:
 			il_func = core.BNGetFlowGraphLowLevelILFunction(self.handle)
@@ -549,6 +554,14 @@ class FlowGraph(object):
 			if function is None:
 				return None
 			return mediumlevelil.MediumLevelILFunction(function.arch, il_func, function)
+		if self.is_high_level_il:
+			il_func = core.BNGetFlowGraphHighLevelILFunction(self.handle)
+			if not il_func:
+				return None
+			function = self.function
+			if function is None:
+				return None
+			return highlevelil.HighLevelILFunction(function.arch, il_func, function)
 		return None
 
 	@il_function.setter
@@ -556,12 +569,19 @@ class FlowGraph(object):
 		if isinstance(func, lowlevelil.LowLevelILFunction):
 			core.BNSetFlowGraphLowLevelILFunction(self.handle, func.handle)
 			core.BNSetFlowGraphMediumLevelILFunction(self.handle, None)
+			core.BNSetFlowGraphHighLevelILFunction(self.handle, None)
 		elif isinstance(func, mediumlevelil.MediumLevelILFunction):
 			core.BNSetFlowGraphLowLevelILFunction(self.handle, None)
 			core.BNSetFlowGraphMediumLevelILFunction(self.handle, func.handle)
+			core.BNSetFlowGraphHighLevelILFunction(self.handle, None)
+		elif isinstance(func, highlevelil.HighLevelILFunction):
+			core.BNSetFlowGraphLowLevelILFunction(self.handle, None)
+			core.BNSetFlowGraphMediumLevelILFunction(self.handle, None)
+			core.BNSetFlowGraphHighLevelILFunction(self.handle, func.handle)
 		elif func is None:
 			core.BNSetFlowGraphLowLevelILFunction(self.handle, None)
 			core.BNSetFlowGraphMediumLevelILFunction(self.handle, None)
+			core.BNSetFlowGraphHighLevelILFunction(self.handle, None)
 		else:
 			raise TypeError("expected IL function for setting il_function property")