Add missing get_basic_block_at() function for IL functions in Python API. Close #3339.

3 files changed, 57 insertions, 0 deletions

diff --git a/python/highlevelil.py b/python/highlevelil.py
index 59434fb0..2acaea74 100644
--- a/python/highlevelil.py
+++ b/python/highlevelil.py
@@ -2214,6 +2214,25 @@ class HighLevelILFunction:
 	def basic_blocks(self) -> 'function.HighLevelILBasicBlockList':
 		return function.HighLevelILBasicBlockList(self)
 
+	def get_basic_block_at(self, index: int) -> Optional['basicblock.BasicBlock']:
+		"""
+		``get_basic_block_at`` returns the BasicBlock at the given HLIL instruction ``index``.
+
+		:param int index: Index of the HLIL instruction of the BasicBlock to retrieve.
+		:Example:
+			>>> current_il_function.get_basic_block_at(current_il_index)
+			<llil block: x86@19-26>
+		"""
+		block = core.BNGetHighLevelILBasicBlockForInstruction(self.handle, index)
+		if not block:
+			return None
+
+		view = None
+		if self._source_function is not None:
+			view = self._source_function.view
+
+		return HighLevelILBasicBlock(block, self, view)
+
 	@property
 	def instructions(self) -> Generator[HighLevelILInstruction, None, None]:
 		"""A generator of hlil instructions of the current function"""
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index 98bf4da6..137a2d77 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -2772,6 +2772,25 @@ class LowLevelILFunction:
 	def basic_blocks(self) -> 'function.LowLevelILBasicBlockList':
 		return function.LowLevelILBasicBlockList(self)
 
+	def get_basic_block_at(self, index: int) -> Optional['basicblock.BasicBlock']:
+		"""
+		``get_basic_block_at`` returns the BasicBlock at the given LLIL instruction ``index``.
+
+		:param int index: Index of the LLIL instruction of the BasicBlock to retrieve.
+		:Example:
+			>>> current_il_function.get_basic_block_at(current_il_index)
+			<llil block: x86@19-26>
+		"""
+		block = core.BNGetLowLevelILBasicBlockForInstruction(self.handle, index)
+		if not block:
+			return None
+
+		view = None
+		if self._source_function is not None:
+			view = self._source_function.view
+
+		return LowLevelILBasicBlock(block, self, view)
+
 	@property
 	def instructions(self) -> Generator['LowLevelILInstruction', None, None]:
 		"""A generator of llil instructions of the current llil function"""
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index d79e2984..07ca7a22 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -2758,6 +2758,25 @@ class MediumLevelILFunction:
 	def basic_blocks(self) -> 'function.MediumLevelILBasicBlockList':
 		return function.MediumLevelILBasicBlockList(self)
 
+	def get_basic_block_at(self, index: int) -> Optional['basicblock.BasicBlock']:
+		"""
+		``get_basic_block_at`` returns the BasicBlock at the given MLIL instruction ``index``.
+
+		:param int index: Index of the MLIL instruction of the BasicBlock to retrieve.
+		:Example:
+			>>> current_il_function.get_basic_block_at(current_il_index)
+			<mlil block: x86@40-60>
+		"""
+		block = core.BNGetMediumLevelILBasicBlockForInstruction(self.handle, index)
+		if not block:
+			return None
+
+		view = None
+		if self._source_function is not None:
+			view = self._source_function.view
+
+		return MediumLevelILBasicBlock(block, self, view)
+
 	@property
 	def instructions(self) -> Generator[MediumLevelILInstruction, None, None]:
 		"""A generator of mlil instructions of the current function"""