Basic blocks have incoming and outgoing edges with basic block references, use core object identity for equality

1 files changed, 42 insertions, 14 deletions

diff --git a/python/basicblock.py b/python/basicblock.py
index f6dbd60d..ebacb311 100644
--- a/python/basicblock.py
+++ b/python/basicblock.py
@@ -29,19 +29,17 @@ import function
 
 
 class BasicBlockEdge(object):
-	def __init__(self, branch_type, target, arch):
+	def __init__(self, branch_type, target):
 		self.type = branch_type
-		if self.type != BranchType.UnresolvedBranch:
-			self.target = target
-			self.arch = arch
+		self.target = target
 
 	def __repr__(self):
 		if self.type == BranchType.UnresolvedBranch:
 			return "<%s>" % BranchType(self.type).name
-		elif self.arch:
-			return "<%s: %s@%#x>" % (self.type, self.arch.name, self.target)
+		elif self.target.arch:
+			return "<%s: %s@%#x>" % (BranchType(self.type).name, self.target.arch.name, self.target.start)
 		else:
-			return "<%s: %#x>" % (self.type, self.target)
+			return "<%s: %#x>" % (BranchType(self.type).name, self.target.start)
 
 
 class BasicBlock(object):
@@ -52,6 +50,16 @@ class BasicBlock(object):
 	def __del__(self):
 		core.BNFreeBasicBlock(self.handle)
 
+	def __eq__(self, value):
+		if not isinstance(value, BasicBlock):
+			return False
+		return ctypes.addressof(self.handle.contents) == ctypes.addressof(value.handle.contents)
+
+	def __ne__(self, value):
+		if not isinstance(value, BasicBlock):
+			return True
+		return ctypes.addressof(self.handle.contents) != ctypes.addressof(value.handle.contents)
+
 	@property
 	def function(self):
 		"""Basic block function (read-only)"""
@@ -84,20 +92,40 @@ class BasicBlock(object):
 		return core.BNGetBasicBlockLength(self.handle)
 
 	@property
+	def index(self):
+		"""Basic block index in list of blocks for the function (read-only)"""
+		return core.BNGetBasicBlockIndex(self.handle)
+
+	@property
 	def outgoing_edges(self):
 		"""List of basic block outgoing edges (read-only)"""
 		count = ctypes.c_ulonglong(0)
 		edges = core.BNGetBasicBlockOutgoingEdges(self.handle, count)
 		result = []
 		for i in xrange(0, count.value):
-			branch_type = edges[i].type
-			target = edges[i].target
-			if edges[i].arch:
-				arch = architecture.Architecture(edges[i].arch)
+			branch_type = BranchType(edges[i].type)
+			if edges[i].target:
+				target = BasicBlock(self.view, core.BNNewBasicBlockReference(edges[i].target))
+			else:
+				target = None
+			result.append(BasicBlockEdge(branch_type, target))
+		core.BNFreeBasicBlockEdgeList(edges, count.value)
+		return result
+
+	@property
+	def incoming_edges(self):
+		"""List of basic block incoming edges (read-only)"""
+		count = ctypes.c_ulonglong(0)
+		edges = core.BNGetBasicBlockIncomingEdges(self.handle, count)
+		result = []
+		for i in xrange(0, count.value):
+			branch_type = BranchType(edges[i].type)
+			if edges[i].target:
+				target = BasicBlock(self.view, core.BNNewBasicBlockReference(edges[i].target))
 			else:
-				arch = None
-			result.append(BasicBlockEdge(branch_type, target, arch))
-		core.BNFreeBasicBlockOutgoingEdgeList(edges)
+				target = None
+			result.append(BasicBlockEdge(branch_type, target))
+		core.BNFreeBasicBlockEdgeList(edges, count.value)
 		return result
 
 	@property