From 309fc943fb48d9ce4f6f7c59f7d66bf0ca69eef1 Mon Sep 17 00:00:00 2001
From: Josh F <josh@vector35.com>
Date: Thu, 6 Oct 2022 18:52:48 -0400
Subject: Add APIs for checking and writing snapshot data

---
 binaryninjaapi.h  |  3 +++
 binaryninjacore.h |  3 +++
 database.cpp      | 19 +++++++++++++++++++
 3 files changed, 25 insertions(+)

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 15850e44..a8f761a0 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1320,6 +1320,7 @@ namespace BinaryNinja {
 		std::string GetName();
 		bool IsAutoSave();
 		bool HasContents();
+		bool HasData();
 		bool HasUndo();
 		Ref<Snapshot> GetFirstParent();
 		std::vector<Ref<Snapshot>> GetParents();
@@ -1340,12 +1341,14 @@ namespace BinaryNinja {
 	  public:
 		Database(BNDatabase* database);
 
+		bool SnapshotHasData(int64_t id);
 		Ref<Snapshot> GetSnapshot(int64_t id);
 		std::vector<Ref<Snapshot>> GetSnapshots();
 		void SetCurrentSnapshot(int64_t id);
 		Ref<Snapshot> GetCurrentSnapshot();
 		int64_t WriteSnapshotData(std::vector<int64_t> parents, Ref<BinaryView> file, const std::string& name,
 		    const Ref<KeyValueStore>& data, bool autoSave, const std::function<bool(size_t, size_t)>& progress);
+		bool StoreDataForSnapshot(int64_t id, const Ref<KeyValueStore>& data, const std::function<bool(size_t, size_t)>& progress);
 		void TrimSnapshot(int64_t id);
 		void RemoveSnapshot(int64_t id);
 
diff --git a/binaryninjacore.h b/binaryninjacore.h
index c15fef12..b601a682 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -3125,6 +3125,8 @@ extern "C"
 	BINARYNINJACOREAPI int64_t BNWriteDatabaseSnapshotData(BNDatabase* database, int64_t* parents, size_t parentCount,
 	    BNBinaryView* file, const char* name, BNKeyValueStore* data, bool autoSave, void* ctxt,
 	    bool (*progress)(void*, size_t, size_t));
+	BINARYNINJACOREAPI bool BNStoreDataForSnapshot(BNDatabase* database, int64_t id, BNKeyValueStore* data,
+	    void* ctxt, bool (*progress)(void*, size_t, size_t));
 	BINARYNINJACOREAPI bool BNTrimDatabaseSnapshot(BNDatabase* database, int64_t id);
 	BINARYNINJACOREAPI bool BNRemoveDatabaseSnapshot(BNDatabase* database, int64_t id);
 	BINARYNINJACOREAPI char** BNGetDatabaseGlobalKeys(BNDatabase* database, size_t* count);
@@ -3136,6 +3138,7 @@ extern "C"
 	BINARYNINJACOREAPI BNFileMetadata* BNGetDatabaseFile(BNDatabase* database);
 	BINARYNINJACOREAPI BNKeyValueStore* BNReadDatabaseAnalysisCache(BNDatabase* database);
 	BINARYNINJACOREAPI bool BNWriteDatabaseAnalysisCache(BNDatabase* database, BNKeyValueStore* val);
+	BINARYNINJACOREAPI bool BNSnapshotHasData(BNDatabase* db, int64_t id);
 
 	// Database snapshots
 	BINARYNINJACOREAPI BNSnapshot* BNNewSnapshotReference(BNSnapshot* snapshot);
diff --git a/database.cpp b/database.cpp
index ba20e02b..e60e0d92 100644
--- a/database.cpp
+++ b/database.cpp
@@ -390,6 +390,19 @@ int64_t Database::WriteSnapshotData(std::vector<int64_t> parents, Ref<BinaryView
 }
 
 
+bool Database::StoreDataForSnapshot(int64_t id, const Ref<KeyValueStore>& data, const std::function<bool(size_t, size_t)>& progress)
+{
+	ProgressContext pctxt;
+	pctxt.callback = progress;
+	bool result = BNStoreDataForSnapshot(m_object, id, data->GetObject(), &pctxt, ProgressCallback);
+	if (!result)
+	{
+		throw DatabaseException("BNStoreDataForSnapshot");
+	}
+	return result;
+}
+
+
 void Database::TrimSnapshot(int64_t id)
 {
 	if (!BNTrimDatabaseSnapshot(m_object, id))
@@ -518,3 +531,9 @@ void Database::WriteAnalysisCache(Ref<KeyValueStore> val)
 		throw DatabaseException("BNWriteDatabaseAnalysisCache");
 	}
 }
+
+
+bool Database::SnapshotHasData(int64_t id)
+{
+	return BNSnapshotHasData(m_object, id);
+}
-- 
cgit v1.3.1