Adding API to retrieve basic block annotations

5 files changed, 65 insertions, 0 deletions

diff --git a/basicblock.cpp b/basicblock.cpp
index e416f7bf..5acf41d7 100644
--- a/basicblock.cpp
+++ b/basicblock.cpp
@@ -90,3 +90,9 @@ void BasicBlock::MarkRecentUse()
 {
 	BNMarkBasicBlockAsRecentlyUsed(m_object);
 }
+
+
+vector<vector<InstructionTextToken>> BasicBlock::GetAnnotations()
+{
+	return GetFunction()->GetBlockAnnotations(GetArchitecture(), GetStart());
+}
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index c9c19953..50def47f 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1275,6 +1275,8 @@ namespace BinaryNinja
 		bool HasUndeterminedOutgoingEdges() const;
 
 		void MarkRecentUse();
+
+		std::vector<std::vector<InstructionTextToken>> GetAnnotations();
 	};
 
 	struct StackVariable
@@ -1385,6 +1387,8 @@ namespace BinaryNinja
 
 		std::vector<IndirectBranchInfo> GetIndirectBranches();
 		std::vector<IndirectBranchInfo> GetIndirectBranchesAt(Architecture* arch, uint64_t addr);
+
+		std::vector<std::vector<InstructionTextToken>> GetBlockAnnotations(Architecture* arch, uint64_t addr);
 	};
 
 	struct FunctionGraphTextLine
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 0eff6aaf..db6f6f27 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -496,6 +496,12 @@ extern "C"
 		uint64_t value;
 	};
 
+	struct BNInstructionTextLine
+	{
+		BNInstructionTextToken* tokens;
+		size_t count;
+	};
+
 	struct BNCustomArchitecture
 	{
 		void* context;
@@ -1179,6 +1185,10 @@ extern "C"
 	                                                                 uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI void BNFreeIndirectBranchList(BNIndirectBranchInfo* branches);
 
+	BINARYNINJACOREAPI BNInstructionTextLine* BNGetFunctionBlockAnnotations(BNFunction* func, BNArchitecture* arch,
+		uint64_t addr, size_t* count);
+	BINARYNINJACOREAPI void BNFreeInstructionTextLines(BNInstructionTextLine* lines, size_t count);
+
 	// Function graph
 	BINARYNINJACOREAPI BNFunctionGraph* BNCreateFunctionGraph(BNFunction* func);
 	BINARYNINJACOREAPI BNFunctionGraph* BNNewFunctionGraphReference(BNFunctionGraph* graph);
diff --git a/function.cpp b/function.cpp
index 7cc30351..70cbfe2a 100644
--- a/function.cpp
+++ b/function.cpp
@@ -458,3 +458,28 @@ vector<IndirectBranchInfo> Function::GetIndirectBranchesAt(Architecture* arch, u
 	BNFreeIndirectBranchList(branches);
 	return result;
 }
+
+
+vector<vector<InstructionTextToken>> Function::GetBlockAnnotations(Architecture* arch, uint64_t addr)
+{
+	size_t count;
+	BNInstructionTextLine* lines = BNGetFunctionBlockAnnotations(m_object, arch->GetObject(), addr, &count);
+
+	vector<vector<InstructionTextToken>> result;
+	for (size_t i = 0; i < count; i++)
+	{
+		vector<InstructionTextToken> line;
+		for (size_t j = 0; j < lines[i].count; j++)
+		{
+			InstructionTextToken token;
+			token.type = lines[i].tokens[j].type;
+			token.text = lines[i].tokens[j].text;
+			token.value = lines[i].tokens[j].value;
+			line.push_back(token);
+		}
+		result.push_back(line);
+	}
+
+	BNFreeInstructionTextLines(lines, count);
+	return result;
+}
diff --git a/python/__init__.py b/python/__init__.py
index 8ddb6262..821e3f41 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -2408,6 +2408,21 @@ class Function(object):
 		core.BNFreeIndirectBranchList(branches)
 		return result
 
+	def get_block_annotations(self, arch, addr):
+		count = ctypes.c_ulonglong(0)
+		lines = core.BNGetFunctionBlockAnnotations(self.handle, arch.handle, addr, count)
+		result = []
+		for i in xrange(0, count.value):
+			tokens = []
+			for j in xrange(0, lines[i].count):
+				token_type = core.BNInstructionTextTokenType_names[lines[i].tokens[j].type]
+				text = lines[i].tokens[j].text
+				value = lines[i].tokens[j].value
+				tokens.append(InstructionTextToken(token_type, text, value))
+			result.append(tokens)
+		core.BNFreeInstructionTextLines(lines, count.value)
+		return result
+
 class BasicBlockEdge:
 	def __init__(self, branch_type, target, arch):
 		self.type = core.BNBranchType_names[branch_type]
@@ -2484,6 +2499,11 @@ class BasicBlock(object):
 		"""Whether basic block has undetermined outgoing edges (read-only)"""
 		return core.BNBasicBlockHasUndeterminedOutgoingEdges(self.handle)
 
+	@property
+	def annotations(self):
+		"""List of automatic annotations for the start of this block (read-only)"""
+		return self.function.get_block_annotations(self.arch, self.start)
+
 	def __setattr__(self, name, value):
 		try:
 			object.__setattr__(self,name,value)