Merginging with dev

1 files changed, 67 insertions, 15 deletions

diff --git a/python/function.py b/python/function.py
index 7d03585c..86c93bf7 100644
--- a/python/function.py
+++ b/python/function.py
@@ -26,7 +26,7 @@ import ctypes
 import _binaryninjacore as core
 from enums import (FunctionGraphType, BranchType, SymbolType, InstructionTextTokenType,
 	HighlightStandardColor, HighlightColorStyle, RegisterValueType, ImplicitRegisterExtend,
-	DisassemblyOption, IntegerDisplayType)
+	DisassemblyOption, IntegerDisplayType, InstructionTextTokenContext)
 import architecture
 import highlight
 import associateddatastore
@@ -177,6 +177,16 @@ class Function(object):
 			core.BNReleaseAdvancedFunctionAnalysisDataMultiple(self.handle, self._advanced_analysis_requests)
 		core.BNFreeFunction(self.handle)
 
+	def __eq__(self, value):
+		if not isinstance(value, Function):
+			return False
+		return ctypes.addressof(self.handle.contents) == ctypes.addressof(value.handle.contents)
+
+	def __ne__(self, value):
+		if not isinstance(value, Function):
+			return True
+		return ctypes.addressof(self.handle.contents) != ctypes.addressof(value.handle.contents)
+
 	@classmethod
 	def _unregister(cls, func):
 		handle = ctypes.cast(func, ctypes.c_void_p)
@@ -669,7 +679,9 @@ class Function(object):
 				value = lines[i].tokens[j].value
 				size = lines[i].tokens[j].size
 				operand = lines[i].tokens[j].operand
-				tokens.append(InstructionTextToken(token_type, text, value, size, operand))
+				context = lines[i].tokens[j].context
+				address = lines[i].tokens[j].address
+				tokens.append(InstructionTextToken(token_type, text, value, size, operand, context, address))
 			result.append(tokens)
 		core.BNFreeInstructionTextLines(lines, count.value)
 		return result
@@ -764,7 +776,9 @@ class Function(object):
 		"""
 		if arch is None:
 			arch = self.arch
-		if not isinstance(color, highlight.HighlightColor):
+		if not isinstance(color, HighlightStandardColor) and not isinstance(color, highlight.HighlightColor):
+			raise ValueError("Specified color is not one of HighlightStandardColor, highlight.HighlightColor")
+		if isinstance(color, HighlightStandardColor):
 			color = highlight.HighlightColor(color = color)
 		core.BNSetAutoInstructionHighlight(self.handle, arch.handle, addr, color._get_core_struct())
 
@@ -784,6 +798,8 @@ class Function(object):
 			arch = self.arch
 		if not isinstance(color, HighlightStandardColor) and not isinstance(color, highlight.HighlightColor):
 			raise ValueError("Specified color is not one of HighlightStandardColor, highlight.HighlightColor")
+		if isinstance(color, HighlightStandardColor):
+			color = highlight.HighlightColor(color)
 		core.BNSetUserInstructionHighlight(self.handle, arch.handle, addr, color._get_core_struct())
 
 
@@ -831,16 +847,19 @@ class DisassemblyTextLine(object):
 
 
 class FunctionGraphEdge(object):
-	def __init__(self, branch_type, arch, target, points):
+	def __init__(self, branch_type, source, target, points):
 		self.type = BranchType(branch_type)
-		self.arch = arch
+		self.source = source
 		self.target = target
 		self.points = points
 
 	def __repr__(self):
-		if self.arch:
-			return "<%s: %s@%#x>" % (self.type.name, self.arch.name, self.target)
-		return "<%s: %#x>" % (self.type, self.target)
+		return "<%s: %s>" % (self.type.name, repr(self.target))
+
+	@property
+	def back_edge(self):
+		"""Whether the edge is a back edge (end of a loop)"""
+		return self.target in self.source.basic_block.dominators
 
 
 class FunctionGraphBlock(object):
@@ -850,6 +869,16 @@ class FunctionGraphBlock(object):
 	def __del__(self):
 		core.BNFreeFunctionGraphBlock(self.handle)
 
+	def __eq__(self, value):
+		if not isinstance(value, FunctionGraphBlock):
+			return False
+		return ctypes.addressof(self.handle.contents) == ctypes.addressof(value.handle.contents)
+
+	def __ne__(self, value):
+		if not isinstance(value, FunctionGraphBlock):
+			return True
+		return ctypes.addressof(self.handle.contents) != ctypes.addressof(value.handle.contents)
+
 	@property
 	def basic_block(self):
 		"""Basic block associated with this part of the function graph (read-only)"""
@@ -916,7 +945,9 @@ class FunctionGraphBlock(object):
 				value = lines[i].tokens[j].value
 				size = lines[i].tokens[j].size
 				operand = lines[i].tokens[j].operand
-				tokens.append(InstructionTextToken(token_type, text, value, size, operand))
+				context = lines[i].tokens[j].context
+				address = lines[i].tokens[j].address
+				tokens.append(InstructionTextToken(token_type, text, value, size, operand, context, address))
 			result.append(DisassemblyTextLine(addr, tokens))
 		core.BNFreeDisassemblyTextLines(lines, count.value)
 		return result
@@ -930,13 +961,19 @@ class FunctionGraphBlock(object):
 		for i in xrange(0, count.value):
 			branch_type = BranchType(edges[i].type)
 			target = edges[i].target
-			arch = None
-			if edges[i].arch is not None:
-				arch = architecture.Architecture(edges[i].arch)
+			if target:
+				func = core.BNGetBasicBlockFunction(target)
+				if func is None:
+					core.BNFreeBasicBlock(target)
+					target = None
+				else:
+					target = basicblock.BasicBlock(binaryview.BinaryView(handle = core.BNGetFunctionData(func)),
+						core.BNNewBasicBlockReference(target))
+					core.BNFreeFunction(func)
 			points = []
 			for j in xrange(0, edges[i].pointCount):
 				points.append((edges[i].points[j].x, edges[i].points[j].y))
-			result.append(FunctionGraphEdge(branch_type, arch, target, points))
+			result.append(FunctionGraphEdge(branch_type, self, target, points))
 		core.BNFreeFunctionGraphBlockOutgoingEdgeList(edges, count.value)
 		return result
 
@@ -966,7 +1003,9 @@ class FunctionGraphBlock(object):
 					value = lines[i].tokens[j].value
 					size = lines[i].tokens[j].size
 					operand = lines[i].tokens[j].operand
-					tokens.append(InstructionTextToken(token_type, text, value, size, operand))
+					context = lines[i].tokens[j].context
+					address = lines[i].tokens[j].address
+					tokens.append(InstructionTextToken(token_type, text, value, size, operand, context, address))
 				yield DisassemblyTextLine(addr, tokens)
 		finally:
 			core.BNFreeDisassemblyTextLines(lines, count.value)
@@ -1020,6 +1059,16 @@ class FunctionGraph(object):
 		self.abort()
 		core.BNFreeFunctionGraph(self.handle)
 
+	def __eq__(self, value):
+		if not isinstance(value, FunctionGraph):
+			return False
+		return ctypes.addressof(self.handle.contents) == ctypes.addressof(value.handle.contents)
+
+	def __ne__(self, value):
+		if not isinstance(value, FunctionGraph):
+			return True
+		return ctypes.addressof(self.handle.contents) != ctypes.addressof(value.handle.contents)
+
 	@property
 	def function(self):
 		"""Function for a function graph (read-only)"""
@@ -1237,12 +1286,15 @@ class InstructionTextToken(object):
 		========================== ============================================
 
 	"""
-	def __init__(self, token_type, text, value = 0, size = 0, operand = 0xffffffff):
+	def __init__(self, token_type, text, value = 0, size = 0, operand = 0xffffffff,
+		context = InstructionTextTokenContext.NoTokenContext, address = 0):
 		self.type = InstructionTextTokenType(token_type)
 		self.text = text
 		self.value = value
 		self.size = size
 		self.operand = operand
+		self.context = InstructionTextTokenContext(context)
+		self.address = address
 
 	def __str__(self):
 		return self.text