From 1f885d1482d0d4e91f84c6d171672bf0d545923e Mon Sep 17 00:00:00 2001
From: rollsafe <rollsafe@users.noreply.github.com>
Date: Tue, 11 Jun 2019 15:29:57 -0400
Subject: Add comments map to binaryview

---
 binaryninjaapi.h     |  4 ++++
 binaryninjacore.h    |  4 ++++
 binaryview.cpp       | 26 ++++++++++++++++++++++++++
 python/binaryview.py | 45 +++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 79 insertions(+)

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 53cacfa8..096f089d 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1541,6 +1541,10 @@ namespace BinaryNinja
 
 		std::vector<std::string> GetUniqueSectionNames(const std::vector<std::string>& names);
 
+		std::string GetCommentForAddress(uint64_t addr) const;
+		std::vector<uint64_t> GetCommentedAddresses() const;
+		void SetCommentForAddress(uint64_t addr, const std::string& comment);
+
 		std::vector<BNAddressRange> GetAllocatedRanges();
 
 		void StoreMetadata(const std::string& key, Ref<Metadata> value);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index b24e675e..903ea473 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -2513,6 +2513,10 @@ extern "C"
 	BINARYNINJACOREAPI BNFunction** BNGetAnalysisFunctionsForAddress(BNBinaryView* view, uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI BNFunction* BNGetAnalysisEntryPoint(BNBinaryView* view);
 
+	BINARYNINJACOREAPI char* BNGetGlobalCommentForAddress(BNBinaryView* view, uint64_t addr);
+	BINARYNINJACOREAPI uint64_t* BNGetGlobalCommentedAddresses(BNBinaryView* view, size_t* count);
+	BINARYNINJACOREAPI void BNSetGlobalCommentForAddress(BNBinaryView* view, uint64_t addr, const char* comment);
+
 	BINARYNINJACOREAPI BNBinaryView* BNGetFunctionData(BNFunction* func);
 	BINARYNINJACOREAPI BNArchitecture* BNGetFunctionArchitecture(BNFunction* func);
 	BINARYNINJACOREAPI BNPlatform* BNGetFunctionPlatform(BNFunction* func);
diff --git a/binaryview.cpp b/binaryview.cpp
index 5811b953..f95bf2fd 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -2333,6 +2333,32 @@ vector<string> BinaryView::GetUniqueSectionNames(const vector<string>& names)
 }
 
 
+string BinaryView::GetCommentForAddress(uint64_t addr) const
+{
+	char* comment = BNGetGlobalCommentForAddress(m_object, addr);
+	string result = comment;
+	BNFreeString(comment);
+	return result;
+}
+
+
+vector<uint64_t> BinaryView::GetCommentedAddresses() const
+{
+	size_t count;
+	uint64_t* addrs = BNGetGlobalCommentedAddresses(m_object, &count);
+	vector<uint64_t> result;
+	result.insert(result.end(), addrs, &addrs[count]);
+	BNFreeAddressList(addrs);
+	return result;
+}
+
+
+void BinaryView::SetCommentForAddress(uint64_t addr, const string& comment)
+{
+	BNSetGlobalCommentForAddress(m_object, addr, comment.c_str());
+}
+
+
 vector<BNAddressRange> BinaryView::GetAllocatedRanges()
 {
 	size_t count;
diff --git a/python/binaryview.py b/python/binaryview.py
index 04f0e2dc..9692689c 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -4373,6 +4373,51 @@ class BinaryView(object):
 		core.BNFreeStringList(outgoing_names, len(name_list))
 		return result
 
+	@property
+	def address_comments(self):
+		"""
+		Returns a read-only dict of the address comments attached to this BinaryView
+
+		Note that these are different from function-level comments which are specific to each function.
+		For annotating code, it is recommended to use comments attached to functions rather than address
+		comments attached to the BinaryView. On the other hand, BinaryView comments can be attached to data
+		whereas function comments cannot.
+		To create a function-level comment, use :func:`~binaryninja.function.Function.add_user_code_ref`.
+		"""
+		count = ctypes.c_ulonglong()
+		addrs = core.BNGetGlobalCommentedAddresses(self.handle, count)
+		result = {}
+		for i in range(0, count.value):
+			result[addrs[i]] = self.get_comment_at(addrs[i])
+		core.BNFreeAddressList(addrs)
+		return result
+
+	def get_comment_at(self, addr):
+		"""
+		``get_comment_at`` returns the address-based comment attached to the given address in this BinaryView
+		Note that address-based comments are different from function-level comments which are specific to each function.
+		For more information, see :func:`address_comments`.
+	
+		"""
+		return core.BNGetGlobalCommentForAddress(self.handle, addr)
+
+	def set_comment_at(self, addr, comment):
+		"""
+		``set_comment_at`` sets a comment for the BinaryView at the address specified
+		
+		Note that these are different from function-level comments which are specific to each function.
+		For more information, see :func:`address_comments`.
+
+		:param addr int: virtual address within the current BinaryView to apply the comment to
+		:param comment str: string comment to apply
+		:rtype: None
+		:Example:
+
+			>>> bv.set_comment_at(here, "hi")
+
+		"""
+		core.BNSetGlobalCommentForAddress(self.handle, addr, comment)
+
 	def query_metadata(self, key):
 		"""
 		`query_metadata` retrieves a metadata associated with the given key stored in the current BinaryView.
-- 
cgit v1.3.1