Add API for list of basic blocks starting at an address

4 files changed, 32 insertions, 0 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 9c7ff4f1..d4df479b 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -867,6 +867,7 @@ namespace BinaryNinja
 
 		Ref<BasicBlock> GetRecentBasicBlockForAddress(uint64_t addr);
 		std::vector<Ref<BasicBlock>> GetBasicBlocksForAddress(uint64_t addr);
+		std::vector<Ref<BasicBlock>> GetBasicBlocksStartingAtAddress(uint64_t addr);
 
 		std::vector<ReferenceSource> GetCodeReferences(uint64_t addr);
 		std::vector<ReferenceSource> GetCodeReferences(uint64_t addr, uint64_t len);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index f917c2fa..51bf3391 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1420,6 +1420,7 @@ extern "C"
 	BINARYNINJACOREAPI void BNFreeBasicBlockList(BNBasicBlock** blocks, size_t count);
 	BINARYNINJACOREAPI BNBasicBlock* BNGetRecentBasicBlockForAddress(BNBinaryView* view, uint64_t addr);
 	BINARYNINJACOREAPI BNBasicBlock** BNGetBasicBlocksForAddress(BNBinaryView* view, uint64_t addr, size_t* count);
+	BINARYNINJACOREAPI BNBasicBlock** BNGetBasicBlocksStartingAtAddress(BNBinaryView* view, uint64_t addr, size_t* count);
 
 	BINARYNINJACOREAPI BNLowLevelILFunction* BNGetFunctionLowLevelIL(BNFunction* func);
 	BINARYNINJACOREAPI size_t BNGetLowLevelILForInstruction(BNFunction* func, BNArchitecture* arch, uint64_t addr);
diff --git a/binaryview.cpp b/binaryview.cpp
index 1c03eca5..0a7f23ce 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -987,6 +987,20 @@ vector<Ref<BasicBlock>> BinaryView::GetBasicBlocksForAddress(uint64_t addr)
 }
 
 
+vector<Ref<BasicBlock>> BinaryView::GetBasicBlocksStartingAtAddress(uint64_t addr)
+{
+	size_t count;
+	BNBasicBlock** blocks = BNGetBasicBlocksStartingAtAddress(m_object, addr, &count);
+
+	vector<Ref<BasicBlock>> result;
+	for (size_t i = 0; i < count; i++)
+		result.push_back(new BasicBlock(BNNewBasicBlockReference(blocks[i])));
+
+	BNFreeBasicBlockList(blocks, count);
+	return result;
+}
+
+
 vector<ReferenceSource> BinaryView::GetCodeReferences(uint64_t addr)
 {
 	size_t count;
diff --git a/python/__init__.py b/python/__init__.py
index 7978fd2c..33202a2c 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -2288,6 +2288,22 @@ class BinaryView(object):
 		core.BNFreeBasicBlockList(blocks, count.value)
 		return result
 
+	def get_basic_blocks_starting_at(self, addr):
+		"""
+		``get_basic_blocks_at`` get a list of :py:Class:`BasicBlock` objects which start at the provided virtual address.
+
+		:param int addr: virtual address of BasicBlock desired
+		:return: a list of :py:Class:`BasicBlock` objects
+		:rtype: list(BasicBlock)
+		"""
+		count = ctypes.c_ulonglong(0)
+		blocks = core.BNGetBasicBlocksStartingAtAddress(self.handle, addr, count)
+		result = []
+		for i in xrange(0, count.value):
+			result.append(BasicBlock(self, core.BNNewBasicBlockReference(blocks[i])))
+		core.BNFreeBasicBlockList(blocks, count.value)
+		return result
+
 	def get_recent_basic_block_at(self, addr):
 		block = core.BNGetRecentBasicBlockForAddress(self.handle, addr)
 		if block is None: