From 3f4b372a6cef322720467d0a8f41557f0570a618 Mon Sep 17 00:00:00 2001
From: Ryan Stortz <ryan@withzombies.com>
Date: Thu, 14 Apr 2016 20:48:04 -0400
Subject: Modifies the python api to allow iteration over some base types

Introduces __iter__ to Function, which yields BasicBlocks
Introduces __iter__ to BasicBlock, with yields instruction text
Extended BasicBlock with ILBasicBlock, to iterate over ILBasicBlocks
(which yields ILInstruction instances)
---
 python/__init__.py | 50 +++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 47 insertions(+), 3 deletions(-)

(limited to 'python')

diff --git a/python/__init__.py b/python/__init__.py
index e1429bd6..1da42f38 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -1932,7 +1932,7 @@ class Function:
 			return result
 		elif name == "low_level_il":
 			return LowLevelILFunction(self.arch, core.BNNewLowLevelILFunctionReference(
-						core.BNGetFunctionLowLevelIL(self.handle)))
+						core.BNGetFunctionLowLevelIL(self.handle)), self)
 		elif name == "low_level_il_basic_blocks":
 			count = ctypes.c_ulonglong()
 			blocks = core.BNGetFunctionLowLevelILBasicBlockList(self.handle, count)
@@ -1954,6 +1954,14 @@ class Function:
 			return result
 		raise AttributeError, "no attribute '%s'" % name
 
+	def __iter__(self):
+		count = ctypes.c_ulonglong()
+		blocks = core.BNGetFunctionBasicBlockList(self.handle, count)
+		for i in xrange(0, count.value):
+			yield BasicBlock(self._view, core.BNNewBasicBlockReference(blocks[i]))
+	    	core.BNFreeBasicBlockList(blocks, count.value)
+		raise StopIteration
+
 	def __setattr__(self, name, value):
 		if ((name == "view") or (name == "arch") or (name == "start") or (name == "symbol") or (name == "auto") or
 		(name == "can_return") or (name == "basic_blocks") or (name == "comments") or (name == "low_level_il") or
@@ -2042,7 +2050,7 @@ class Function:
 		result = []
 		for i in xrange(0, count.value):
 			result.append(StackVariableReference(refs[i].sourceOperand, Type(core.BNNewTypeReference(refs[i].type)),
-			                                     refs[i].name, refs[i].startingOffset, refs[i].referencedOffset))
+							     refs[i].name, refs[i].startingOffset, refs[i].referencedOffset))
 		core.BNFreeStackVariableReferenceList(refs, count.value)
 		return result
 
@@ -2131,9 +2139,36 @@ class BasicBlock:
 		else:
 			return "<block: 0x%x-0x%x>" % (self.start, self.end)
 
+	def __iter__(self):
+		func = getattr(self, "function")
+		start = getattr(self, "start")
+		end = getattr(self, "end")
+
+		idx = start
+		while idx < end:
+			data = self.view.read(idx, 16)
+			inst_info = self.view.arch.get_instruction_info(data, idx)
+			inst_text = self.view.arch.get_instruction_text(data, idx)
+			
+			yield inst_text
+			idx += inst_info.length
+		raise StopIteration
+
 	def mark_recent_use():
 		core.BNMarkBasicBlockAsRecentlyUsed(self.handle)
 
+class ILBasicBlock(BasicBlock):
+	def __iter__(self):
+		func = getattr(self, "function")
+		low_level_il_function = func.low_level_il
+		start = getattr(self, "start")
+		end = getattr(self, "end")
+
+		for idx in xrange(start, end):
+			yield low_level_il_function[idx]
+
+		raise StopIteration
+
 class FunctionGraphTextLine:
 	def __init__(self, addr, tokens):
 		self.address = addr
@@ -3390,8 +3425,9 @@ class LowLevelILExpr:
 		self.index = index
 
 class LowLevelILFunction:
-	def __init__(self, arch, handle = None):
+	def __init__(self, arch, handle = None, mcfunc = None):
 		self.arch = arch
+		self.mcfunc = mcfunc
 		if handle is not None:
 			self.handle = core.handle_of_type(handle, core.BNLowLevelILFunction)
 		else:
@@ -3429,6 +3465,14 @@ class LowLevelILFunction:
 	def __setitem__(self, i):
 		raise IndexError, "instruction modification not implemented"
 
+	def __iter__(self):
+	    count = ctypes.c_ulonglong()
+	    blocks = core.BNGetFunctionLowLevelILBasicBlockList(self.mcfunc.handle, count)
+	    for i in xrange(0, count.value):
+		yield ILBasicBlock(self.mcfunc._view, core.BNNewBasicBlockReference(blocks[i]))
+	    core.BNFreeBasicBlockList(blocks, count.value)
+	    raise StopIteration
+
 	def expr(self, operation, a = 0, b = 0, c = 0, d = 0, size = 0, flags = None):
 		if isinstance(operation, str):
 			operation = BNLowLevelILOperation_by_name[operation]
-- 
cgit v1.3.1