Expose some BasicBlock APIs.

4 files changed, 33 insertions, 1 deletions

diff --git a/basicblock.cpp b/basicblock.cpp
index 3e3fccef..1675afdf 100644
--- a/basicblock.cpp
+++ b/basicblock.cpp
@@ -200,6 +200,12 @@ bool BasicBlock::CanExit() const
 }
 
 
+void BasicBlock::SetCanExit(bool value)
+{
+	BNBasicBlockSetCanExit(m_object, value);
+}
+
+
 set<Ref<BasicBlock>> BasicBlock::GetDominators(bool post) const
 {
 	size_t count;
@@ -481,3 +487,12 @@ bool BasicBlock::GetInstructionContainingAddress(uint64_t addr, uint64_t* start)
 {
 	return BNGetBasicBlockInstructionContainingAddress(m_object, addr, start);
 }
+
+
+Ref<BasicBlock> BasicBlock::GetSourceBlock() const
+{
+	BNBasicBlock* block = BNGetBasicBlockSourceBlock(m_object);
+	if (!block)
+		return nullptr;
+	return new BasicBlock(block);
+}
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index ea67f459..9f7acf0f 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -3132,6 +3132,7 @@ __attribute__ ((format (printf, 1, 2)))
 		std::vector<BasicBlockEdge> GetIncomingEdges() const;
 		bool HasUndeterminedOutgoingEdges() const;
 		bool CanExit() const;
+		void SetCanExit(bool value);
 
 		std::set<Ref<BasicBlock>> GetDominators(bool post = false) const;
 		std::set<Ref<BasicBlock>> GetStrictDominators(bool post = false) const;
@@ -3167,6 +3168,7 @@ __attribute__ ((format (printf, 1, 2)))
 		Ref<MediumLevelILFunction> GetMediumLevelILFunction() const;
 
 		bool GetInstructionContainingAddress(uint64_t addr, uint64_t* start);
+		Ref<BasicBlock> GetSourceBlock() const;
 	};
 
 	struct VariableNameAndType
diff --git a/binaryninjacore.h b/binaryninjacore.h
index c954b445..eef37ab6 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -3448,6 +3448,7 @@ __attribute__ ((format (printf, 1, 2)))
 	BINARYNINJACOREAPI void BNFreeBasicBlockEdgeList(BNBasicBlockEdge* edges, size_t count);
 	BINARYNINJACOREAPI bool BNBasicBlockHasUndeterminedOutgoingEdges(BNBasicBlock* block);
 	BINARYNINJACOREAPI bool BNBasicBlockCanExit(BNBasicBlock* block);
+	BINARYNINJACOREAPI void BNBasicBlockSetCanExit(BNBasicBlock* block, bool value);
 	BINARYNINJACOREAPI bool BNBasicBlockHasInvalidInstructions(BNBasicBlock* block);
 	BINARYNINJACOREAPI size_t BNGetBasicBlockIndex(BNBasicBlock* block);
 	BINARYNINJACOREAPI BNBasicBlock** BNGetBasicBlockDominators(BNBasicBlock* block, size_t* count, bool post);
@@ -3465,6 +3466,7 @@ __attribute__ ((format (printf, 1, 2)))
 	BINARYNINJACOREAPI BNMediumLevelILFunction* BNGetBasicBlockMediumLevelILFunction(BNBasicBlock* block);
 	BINARYNINJACOREAPI bool BNGetBasicBlockInstructionContainingAddress(BNBasicBlock* block,
 		uint64_t addr, uint64_t* start);
+	BINARYNINJACOREAPI BNBasicBlock* BNGetBasicBlockSourceBlock(BNBasicBlock* block);
 
 	BINARYNINJACOREAPI BNDisassemblyTextLine* BNGetBasicBlockDisassemblyText(BNBasicBlock* block,
 		BNDisassemblySettings* settings, size_t* count);
diff --git a/python/basicblock.py b/python/basicblock.py
index 4b1633a1..2394f25f 100644
--- a/python/basicblock.py
+++ b/python/basicblock.py
@@ -297,9 +297,14 @@ class BasicBlock(object):
 
 	@property
 	def can_exit(self):
-		"""Whether basic block can return or is tagged as 'No Return' (read-only)"""
+		"""Whether basic block can return or is tagged as 'No Return'"""
 		return core.BNBasicBlockCanExit(self.handle)
 
+	@can_exit.setter
+	def can_exit(self, value):
+		"""Sets whether basic block can return or is tagged as 'No Return'"""
+		BNBasicBlockSetCanExit(self.handle, value)
+
 	@property
 	def has_invalid_instructions(self):
 		"""Whether basic block has any invalid instructions (read-only)"""
@@ -526,3 +531,11 @@ class BasicBlock(object):
 		start = ctypes.c_uint64()
 		ret = core.BNGetBasicBlockInstructionContainingAddress(self.handle, addr, start)
 		return ret, start.value
+
+	@property
+	def source_block(self):
+		"""Source of this basic block (read-only)"""
+		result = core.BNGetBasicBlockSourceBlock(self.handle)
+		if not result:
+			return None
+		return self._create_instance(result, self.view)