Add ability to index into functions and basicblocks

2 files changed, 99 insertions, 53 deletions

diff --git a/python/basicblock.py b/python/basicblock.py
index 4bf2f3b2..27eadd51 100644
--- a/python/basicblock.py
+++ b/python/basicblock.py
@@ -25,6 +25,7 @@ import binaryninja
 from binaryninja import highlight
 from binaryninja import _binaryninjacore as core
 from binaryninja.enums import BranchType, HighlightColorStyle, HighlightStandardColor, InstructionTextTokenType
+from binaryninja import log
 
 # 2-3 compatibility
 from binaryninja import range
@@ -61,6 +62,8 @@ class BasicBlock(object):
 		self.handle = core.handle_of_type(handle, core.BNBasicBlock)
 		self._arch = None
 		self._func = None
+		self._instStarts = None
+		self._instLengths = None
 
 	def __del__(self):
 		core.BNFreeBasicBlock(self.handle)
@@ -75,12 +78,71 @@ class BasicBlock(object):
 			return True
 		return ctypes.addressof(self.handle.contents) != ctypes.addressof(value.handle.contents)
 
+	def __setattr__(self, name, value):
+		try:
+			object.__setattr__(self, name, value)
+		except AttributeError:
+			raise AttributeError("attribute '%s' is read only" % name)
+
+	def __len__(self):
+		return int(core.BNGetBasicBlockLength(self.handle))
+
+	def __repr__(self):
+		arch = self.arch
+		if arch:
+			return "<block: %s@%#x-%#x>" % (arch.name, self.start, self.end)
+		else:
+			return "<block: %#x-%#x>" % (self.start, self.end)
+
+	def __iter__(self):
+		if self._instStarts is None:
+			# no cache is instruction start cache is built so we don't build one
+			# because the user is likely iterating only
+			idx = self.start
+			while idx < self.end:
+				data = self.view.read(idx, min(self.arch.max_instr_length, self.end - idx))
+				inst_text = self.arch.get_instruction_text(data, idx)
+				if inst_text[1] == 0:
+					break
+				yield inst_text
+				idx += inst_text[1]
+		else:
+			for start, length in zip(self._instStarts, self._instLengths):
+				inst_text = self.arch.get_instruction_text(self.view.read(start, length), start)
+				if inst_text[1] == 0:
+					break
+				yield inst_text
+
+	def __getitem__(self, i):
+		self._buildStartCache()
+		start = self._instStarts[i]
+		length = self._instLengths[i]
+		data = self.view.read(start, length)
+		return self.arch.get_instruction_text(data, start)
+
+	def __hash__(self):
+		return hash((self.start, self.end, self.arch.name))
+
+	def _buildStartCache(self):
+		if self._instStarts is None:
+			# build the instruction start cache
+			self._instLengths = []
+			self._instStarts = []
+			start = self.start
+			while start < self.end:
+				length = self.view.get_instruction_length(start)
+				self._instLengths.append(length)
+				self._instStarts.append(start)
+				start += length
+
 	def _create_instance(self, handle, view):
 		"""Internal method used to instantiate child instances"""
 		return BasicBlock(handle, view)
 
-	def __hash__(self):
-		return hash((self.start, self.end, self.arch.name))
+	@property
+	def instruction_count(self):
+		self._buildStartCache()
+		return len(self._instStarts)
 
 	@property
 	def function(self):
@@ -290,35 +352,6 @@ class BasicBlock(object):
 		core.BNFreeBasicBlockList(out_blocks, count.value)
 		return result
 
-	def __setattr__(self, name, value):
-		try:
-			object.__setattr__(self, name, value)
-		except AttributeError:
-			raise AttributeError("attribute '%s' is read only" % name)
-
-	def __len__(self):
-		return int(core.BNGetBasicBlockLength(self.handle))
-
-	def __repr__(self):
-		arch = self.arch
-		if arch:
-			return "<block: %s@%#x-%#x>" % (arch.name, self.start, self.end)
-		else:
-			return "<block: %#x-%#x>" % (self.start, self.end)
-
-	def __iter__(self):
-		start = self.start
-		end = self.end
-
-		idx = start
-		while idx < end:
-			data = self.view.read(idx, min(self.arch.max_instr_length, end - idx))
-			inst_text = self.arch.get_instruction_text(data, idx)
-			if inst_text[1] == 0:
-				break
-			yield inst_text
-			idx += inst_text[1]
-
 	def mark_recent_use(self):
 		core.BNMarkBasicBlockAsRecentlyUsed(self.handle)
 
diff --git a/python/function.py b/python/function.py
index b7c9229c..14e140c1 100644
--- a/python/function.py
+++ b/python/function.py
@@ -389,6 +389,41 @@ class Function(object):
 	def __hash__(self):
 		return hash((self.start, self.arch.name, self.platform.name))
 
+	def __getitem__(self, i):
+		count = ctypes.c_ulonglong()
+		blocks = core.BNGetFunctionBasicBlockList(self.handle, count)
+		try:
+			if i < 0:
+				i = count.value + i
+			if i < 0 or i >= count.value:
+				raise IndexError("index out of range")
+			block = binaryninja.basicblock.BasicBlock(core.BNNewBasicBlockReference(blocks[i]), self._view)
+			return block
+		finally:
+			core.BNFreeBasicBlockList(blocks, count.value)
+
+	def __iter__(self):
+		count = ctypes.c_ulonglong()
+		blocks = core.BNGetFunctionBasicBlockList(self.handle, count)
+		try:
+			for i in range(0, count.value):
+				yield binaryninja.basicblock.BasicBlock(core.BNNewBasicBlockReference(blocks[i]), self._view)
+		finally:
+			core.BNFreeBasicBlockList(blocks, count.value)
+
+	def __setattr__(self, name, value):
+		try:
+			object.__setattr__(self, name, value)
+		except AttributeError:
+			raise AttributeError("attribute '%s' is read only" % name)
+
+	def __repr__(self):
+		arch = self.arch
+		if arch:
+			return "<func: %s@%#x>" % (arch.name, self.start)
+		else:
+			return "<func: %#x>" % self.start
+
 	@classmethod
 	def _unregister(cls, func):
 		handle = ctypes.cast(func, ctypes.c_void_p)
@@ -870,28 +905,6 @@ class Function(object):
 			return None
 		return binaryninja.flowgraph.CoreFlowGraph(graph)
 
-	def __iter__(self):
-		count = ctypes.c_ulonglong()
-		blocks = core.BNGetFunctionBasicBlockList(self.handle, count)
-		try:
-			for i in range(0, count.value):
-				yield binaryninja.basicblock.BasicBlock(core.BNNewBasicBlockReference(blocks[i]), self._view)
-		finally:
-			core.BNFreeBasicBlockList(blocks, count.value)
-
-	def __setattr__(self, name, value):
-		try:
-			object.__setattr__(self, name, value)
-		except AttributeError:
-			raise AttributeError("attribute '%s' is read only" % name)
-
-	def __repr__(self):
-		arch = self.arch
-		if arch:
-			return "<func: %s@%#x>" % (arch.name, self.start)
-		else:
-			return "<func: %#x>" % self.start
-
 	def mark_recent_use(self):
 		core.BNMarkFunctionAsRecentlyUsed(self.handle)