Adding Function level comment APIs

4 files changed, 43 insertions, 3 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 1b02e2f1..1cd65744 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -2153,8 +2153,10 @@ namespace BinaryNinja
 		Ref<BasicBlock> GetBasicBlockAtAddress(Architecture* arch, uint64_t addr) const;
 		void MarkRecentUse();
 
+		std::string GetComment() const;
 		std::string GetCommentForAddress(uint64_t addr) const;
 		std::vector<uint64_t> GetCommentedAddresses() const;
+		void SetComment(const std::string& comment);
 		void SetCommentForAddress(uint64_t addr, const std::string& comment);
 
 		Ref<LowLevelILFunction> GetLowLevelIL() const;
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 44858e6f..f81b821c 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -2041,9 +2041,11 @@ extern "C"
 	BINARYNINJACOREAPI void BNSetFunctionAutoType(BNFunction* func, BNType* type);
 	BINARYNINJACOREAPI void BNSetFunctionUserType(BNFunction* func, BNType* type);
 
+	BINARYNINJACOREAPI char* BNGetFunctionComment(BNFunction* func);
 	BINARYNINJACOREAPI char* BNGetCommentForAddress(BNFunction* func, uint64_t addr);
 	BINARYNINJACOREAPI uint64_t* BNGetCommentedAddresses(BNFunction* func, size_t* count);
 	BINARYNINJACOREAPI void BNFreeAddressList(uint64_t* addrs);
+	BINARYNINJACOREAPI void BNSetFunctionComment(BNFunction* func, const char* comment);
 	BINARYNINJACOREAPI void BNSetCommentForAddress(BNFunction* func, uint64_t addr, const char* comment);
 
 	BINARYNINJACOREAPI BNBasicBlock* BNNewBasicBlockReference(BNBasicBlock* block);
diff --git a/function.cpp b/function.cpp
index 550691d0..2d8db04c 100644
--- a/function.cpp
+++ b/function.cpp
@@ -197,6 +197,15 @@ void Function::MarkRecentUse()
 }
 
 
+string Function::GetComment() const
+{
+	char* comment = BNGetFunctionComment(m_object);
+	string result = comment;
+	BNFreeString(comment);
+	return result;
+}
+
+
 string Function::GetCommentForAddress(uint64_t addr) const
 {
 	char* comment = BNGetCommentForAddress(m_object, addr);
@@ -217,6 +226,12 @@ vector<uint64_t> Function::GetCommentedAddresses() const
 }
 
 
+void Function::SetComment(const string& comment)
+{
+	BNSetFunctionComment(m_object, comment.c_str());
+}
+
+
 void Function::SetCommentForAddress(uint64_t addr, const string& comment)
 {
 	BNSetCommentForAddress(m_object, addr, comment.c_str());
diff --git a/python/function.py b/python/function.py
index 4af95738..5daa7b2a 100644
--- a/python/function.py
+++ b/python/function.py
@@ -656,6 +656,16 @@ class Function(object):
 		result = core.BNGetFunctionGlobalPointerValue(self.handle)
 		return RegisterValue(self.arch, result.value, confidence = result.confidence)
 
+	@property
+	def comment(self):
+		"""Gets the comment for the current function"""
+		return core.BNGetFunctionComment(self.handle)
+
+	@comment.setter
+	def comment(self, comment):
+		"""Sets a comment for the current function"""
+		return core.BNSetFunctionComment(self.handle, comment)
+
 	def __iter__(self):
 		count = ctypes.c_ulonglong()
 		blocks = core.BNGetFunctionBasicBlockList(self.handle, count)
@@ -684,11 +694,22 @@ class Function(object):
 	def get_comment_at(self, addr):
 		return core.BNGetCommentForAddress(self.handle, addr)
 
-	def set_comment_at(self, addr, comment):
+	def set_comment(self, addr, comment):
+		"""Deprecated use set_comment_at instead"""
 		core.BNSetCommentForAddress(self.handle, addr, comment)
 
-	def set_comment(self, addr, comment):
-		"""Deprecated"""
+	def set_comment_at(self, addr, comment):
+		"""
+		``set_comment_at`` sets a comment for the current function at the address specified
+
+		:param addr int: virtual address within the current function to apply the comment to
+		:param comment str: string comment to apply
+		:rtype: None
+		:Example:
+
+			>>> current_function.set_comment_at(here, "hi")
+
+		"""
 		core.BNSetCommentForAddress(self.handle, addr, comment)
 
 	def get_low_level_il_at(self, addr, arch=None):