Database python api

4 files changed, 351 insertions, 0 deletions

diff --git a/python/__init__.py b/python/__init__.py
index 09701f46..6b16a463 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -66,6 +66,7 @@ from .variable import *
 from .websocketprovider import *
 from .workflow import *
 from .commonil import *
+from .database import *
 
 
 def shutdown():
diff --git a/python/database.py b/python/database.py
new file mode 100644
index 00000000..bf1a317e
--- /dev/null
+++ b/python/database.py
@@ -0,0 +1,335 @@
+# coding=utf-8
+# Copyright (c) 2015-2021 Vector 35 Inc
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+import ctypes
+from typing import List, Optional, Dict, Callable
+
+import binaryninja
+from . import _binaryninjacore as core
+from . import databuffer
+from . import filemetadata
+
+
+class KeyValueStore:
+    """
+    ``class KeyValueStore`` maintains access to the raw data stored in Snapshots and various
+    other Database-related structures.
+    """
+    def __init__(self, buffer: Optional[databuffer.DataBuffer] = None, handle=None):
+        if handle is not None:
+            self.handle = core.handle_of_type(handle, core.BNKeyValueStore)
+        else:
+            if buffer is None:
+                _handle = core.BNCreateKeyValueStore()
+            else:
+                _handle = core.BNCreateKeyValueStoreFromDataBuffer(buffer.handle)
+            assert _handle is not None
+            self.handle = _handle
+
+    def __del__(self):
+        core.BNFreeKeyValueStore(self.handle)
+
+    def __getitem__(self, item: str) -> databuffer.DataBuffer:
+        return self.get_value(item)
+
+    def __setitem__(self, key: str, value: databuffer.DataBuffer):
+        return self.set_value(key, value)
+
+    @property
+    def keys(self):
+        """Get a list of all keys stored in the kvs (read-only)"""
+        count = ctypes.c_ulonglong(0)
+        value = core.BNGetKeyValueStoreKeys(self.handle, count)
+        assert value is not None
+
+        result = []
+        try:
+            for i in range(0, count.value):
+                result.append(value[i])
+            return result
+        finally:
+            core.BNFreeStringList(value, count)
+
+    def get_value(self, key: str) -> databuffer.DataBuffer:
+        """Get the value for a single key"""
+        handle = core.BNGetKeyValueStoreBuffer(self.handle, key)
+        assert handle is not None
+        return databuffer.DataBuffer(handle=handle)
+
+    def set_value(self, key: str, value: databuffer.DataBuffer):
+        """Set the value for a single key"""
+        core.BNSetKeyValueStoreBuffer(self.handle, key, value.handle)
+
+    @property
+    def serialized_data(self) -> databuffer.DataBuffer:
+        """Get the stored representation of the kvs (read-only)"""
+        handle = core.BNGetKeyValueStoreSerializedData(self.handle)
+        assert handle is not None
+        return databuffer.DataBuffer(handle=handle)
+
+    def begin_namespace(self, name: str):
+        """Begin storing new keys into a namespace"""
+        core.BNBeginKeyValueStoreNamespace(self.handle, name)
+
+    def end_namespace(self):
+        """End storing new keys into a namespace"""
+        core.BNEndKeyValueStoreNamespace(self.handle)
+
+    @property
+    def empty(self) -> bool:
+        """If the kvs is empty (read-only)"""
+        return core.BNIsKeyValueStoreEmpty(self.handle)
+
+    @property
+    def value_size(self) -> bool:
+        """Number of values in the kvs (read-only)"""
+        return core.BNGetKeyValueStoreValueSize(self.handle)
+
+    @property
+    def data_size(self) -> bool:
+        """Length of serialized data (read-only)"""
+        return core.BNGetKeyValueStoreDataSize(self.handle)
+
+    @property
+    def value_storage_size(self) -> bool:
+        """Size of all data in storage (read-only)"""
+        return core.BNGetKeyValueStoreValueStorageSize(self.handle)
+
+    @property
+    def namespace_size(self) -> bool:
+        """Number of namespaces pushed with begin_namespace (read-only)"""
+        return core.BNGetKeyValueStoreNamespaceSize(self.handle)
+
+
+class Snapshot:
+    """
+    ``class Snapshot`` is a model of an individual database snapshot, created on save.
+    """
+    def __init__(self, handle):
+        self.handle = core.handle_of_type(handle, core.BNSnapshot)
+
+    def __del__(self):
+        core.BNFreeSnapshot(self.handle)
+
+    @property
+    def database(self) -> 'Database':
+        """Get the owning database (read-only)"""
+        return Database(handle=core.BNGetSnapshotDatabase(self.handle))
+
+    @property
+    def id(self) -> int:
+        """Get the numerical id (read-only)"""
+        return core.BNGetSnapshotId(self.handle)
+
+    @property
+    def name(self) -> str:
+        """Get the displayed snapshot name (read-only)"""
+        return core.BNGetSnapshotName(self.handle)
+
+    @property
+    def is_auto_save(self) -> bool:
+        """If the snapshot was the result of an auto-save (read-only)"""
+        return core.BNIsSnapshotAutoSave(self.handle)
+
+    @property
+    def first_parent(self) -> Optional['Snapshot']:
+        """Get the first parent of the snapshot, or None if it has no parents (read-only)"""
+        handle = core.BNGetSnapshotFirstParent(self.handle)
+        if handle is None:
+            return None
+        return Snapshot(handle=handle)
+
+    @property
+    def parents(self) -> List['Snapshot']:
+        """Get a list of all parent snapshots of the snapshot (read-only)"""
+        count = ctypes.c_ulonglong(0)
+        parents = core.BNGetSnapshotParents(self.handle, count)
+
+        result = []
+        try:
+            for i in range(0, count.value):
+                handle = core.BNNewSnapshotReference(parents[i])
+                result.append(Snapshot(handle=handle))
+            return result
+        finally:
+            core.BNFreeSnapshotList(parents, count)
+
+    @property
+    def children(self) -> List['Snapshot']:
+        """Get a list of all child snapshots of the snapshot (read-only)"""
+        count = ctypes.c_ulonglong(0)
+        children = core.BNGetSnapshotChildren(self.handle, count)
+
+        result = []
+        try:
+            for i in range(0, count.value):
+                handle = core.BNNewSnapshotReference(children[i])
+                result.append(Snapshot(handle=handle))
+            return result
+        finally:
+            core.BNFreeSnapshotList(children, count)
+
+    @property
+    def file_contents(self) -> databuffer.DataBuffer:
+        """Get a buffer of the raw data at the time of the snapshot (read-only)"""
+        handle = core.BNGetSnapshotFileContents(self.handle)
+        assert handle is not None
+        return databuffer.DataBuffer(handle=handle)
+
+    @property
+    def file_contents_hash(self) -> databuffer.DataBuffer:
+        """Get a hash of the data at the time of the snapshot (read-only)"""
+        handle = core.BNGetSnapshotFileContentsHash(self.handle)
+        assert handle is not None
+        return databuffer.DataBuffer(handle=handle)
+
+    @property
+    def undo_entries(self):
+        """Get a list of undo entries at the time of the snapshot (read-only)"""
+        raise NotImplementedError("GetUndoEntries is not implemented in python")
+
+    @property
+    def data(self) -> KeyValueStore:
+        """Get the backing kvs data with snapshot fields (read-only)"""
+        handle = core.BNReadSnapshotData(self.handle)
+        assert handle is not None
+        return KeyValueStore(handle=handle)
+
+    def has_ancestor(self, other: 'Snapshot') -> bool:
+        """Determine if this snapshot has another as an ancestor"""
+        return core.BNSnapshotHasAncestor(self.handle, other.handle)
+
+
+class Database:
+    """
+    ``class Database`` provides lower level access to raw snapshot data used to construct analysis data
+    """
+    def __init__(self, handle):
+        self.handle = core.handle_of_type(handle, core.BNDatabase)
+
+    def __del__(self):
+        core.BNFreeDatabase(self.handle)
+
+    def __getitem__(self, item: int) -> Optional[Snapshot]:
+        return self.get_snapshot(item)
+
+    def get_snapshot(self, id: int) -> Optional[Snapshot]:
+        """Get a snapshot by its id, or None if no snapshot with that id exists"""
+        snap = core.BNGetDatabaseSnapshot(self.handle, id)
+        if snap is None:
+            return None
+        return Snapshot(handle=snap)
+
+    @property
+    def snapshots(self) -> List[Snapshot]:
+        """Get a list of all snapshots in the database (read-only)"""
+        count = ctypes.c_ulonglong(0)
+        snapshots = core.BNGetDatabaseSnapshots(self.handle, count)
+        assert snapshots is not None
+
+        result = []
+        try:
+            for i in range(0, count.value):
+                handle = core.BNNewSnapshotReference(snapshots[i])
+                result.append(Snapshot(handle=handle))
+            return result
+        finally:
+            core.BNFreeSnapshotList(snapshots, count)
+
+    @property
+    def current_snapshot(self) -> Optional[Snapshot]:
+        """Get the current snapshot"""
+        snap = core.BNGetDatabaseCurrentSnapshot(self.handle)
+        if snap is None:
+            return None
+        return Snapshot(handle=snap)
+
+    @current_snapshot.setter
+    def current_snapshot(self, value: Snapshot):
+        core.BNSetDatabaseCurrentSnapshot(self.handle, value.id)
+
+    def remove_snapsot(self, id: int):
+        """Remove a snapshot in the database by id"""
+        core.BNRemoveDatabaseSnapshot(self.handle, id)
+
+    @property
+    def global_keys(self) -> List[str]:
+        """Get a list of keys for all globals in the database (read-only)"""
+        count = ctypes.c_ulonglong(0)
+        value = core.BNGetDatabaseGlobalKeys(self.handle, count)
+        assert value is not None
+
+        result = []
+        try:
+            for i in range(0, count.value):
+                result.append(value[i])
+            return result
+        finally:
+            core.BNFreeStringList(value, count)
+
+    @property
+    def globals(self) -> Dict[str, str]:
+        """Get a dictionary of all globals (read-only)"""
+        count = ctypes.c_ulonglong(0)
+        value = core.BNGetDatabaseGlobalKeys(self.handle, count)
+        assert value is not None
+
+        result = {}
+        try:
+            for i in range(0, count.value):
+                key = value[i]
+                result[key] = self.read_global(key)
+            return result
+        finally:
+            core.BNFreeStringList(value, count)
+
+    def read_global(self, key: str) -> str:
+        """Get a specific global by key"""
+        value = core.BNReadDatabaseGlobal(self.handle, key)
+        assert value is not None
+        return value
+
+    def write_global(self, key: str, value: str):
+        """Write a global into the database"""
+        core.BNWriteDatabaseGlobal(self.handle, key, value)
+
+    def read_global_data(self, key: str) -> databuffer.DataBuffer:
+        """Get a specific global by key, as a binary buffer"""
+        handle = core.BNReadDatabaseGlobalData(self.handle, key)
+        assert handle is not None
+        return databuffer.DataBuffer(handle=handle)
+
+    def write_global_data(self, key: str, value: databuffer.DataBuffer):
+        """Write a binary buffer into a global in the database"""
+        core.BNWriteDatabaseGlobalData(self.handle, key, value.handle)
+
+    @property
+    def file(self) -> 'filemetadata.FileMetadata':
+        """Get the owning FileMetadata (read-only)"""
+        handle = core.BNGetDatabaseFile(self.handle)
+        assert handle is not None
+        return filemetadata.FileMetadata(handle=handle)
+
+    @property
+    def analysis_cache(self) -> KeyValueStore:
+        """Get the backing analysis cache kvs (read-only)"""
+        handle = core.BNReadDatabaseAnalysisCache(self.handle)
+        assert handle is not None
+        return KeyValueStore(handle=handle)
diff --git a/python/databuffer.py b/python/databuffer.py
index baf35b17..cc40c5ad 100644
--- a/python/databuffer.py
+++ b/python/databuffer.py
@@ -129,6 +129,12 @@ class DataBuffer:
 		ctypes.memmove(buf, data, len(self))
 		return buf.raw
 
+	def __eq__(self, other: 'DataBuffer') -> bool:
+		# Not cryptographically secure
+		if len(self) != len(other):
+			return False
+		return bytes(self) == bytes(other)
+
 	def escape(self) -> str:
 		return core.BNDataBufferToEscapedString(self.handle)
 
diff --git a/python/filemetadata.py b/python/filemetadata.py
index a8c738c2..0f74aa19 100644
--- a/python/filemetadata.py
+++ b/python/filemetadata.py
@@ -29,6 +29,7 @@ from .enums import SaveOption
 from . import associateddatastore #required for _FileMetadataAssociatedDataStore
 from .log import log_error
 from . import binaryview
+from . import database
 
 ProgressFuncType = Callable[[int, int], bool]
 ViewName = str
@@ -245,6 +246,14 @@ class FileMetadata:
 		return binaryview.BinaryView(file_metadata = self, handle = view)
 
 	@property
+	def database(self) -> Optional['database.Database']:
+		"""Gets the backing Database of the file"""
+		handle = core.BNGetFileMetadataDatabase(self.handle)
+		if handle is None:
+			return None
+		return database.Database(handle=handle)
+
+	@property
 	def saved(self) -> bool:
 		"""Boolean result of whether the file has been saved (Inverse of 'modified' property) (read/write)"""
 		return not core.BNIsFileModified(self.handle)