Skip to main content

binaryninja/
database.rs

1pub mod kvs;
2pub mod snapshot;
3pub mod undo;
4
5use binaryninjacore_sys::*;
6use std::collections::HashMap;
7use std::ffi::c_void;
8use std::fmt::Debug;
9use std::path::Path;
10use std::ptr::NonNull;
11
12use crate::binary_view::BinaryView;
13use crate::data_buffer::DataBuffer;
14use crate::database::kvs::KeyValueStore;
15use crate::database::snapshot::{Snapshot, SnapshotId};
16use crate::progress::{NoProgressCallback, ProgressCallback};
17use crate::rc::{Array, Ref, RefCountable};
18use crate::string::{BnString, IntoCStr};
19
20pub struct Database {
21    pub(crate) handle: NonNull<BNDatabase>,
22}
23
24impl Database {
25    pub(crate) unsafe fn ref_from_raw(handle: NonNull<BNDatabase>) -> Ref<Self> {
26        Ref::new(Self { handle })
27    }
28
29    /// Open a database with the given file path
30    pub fn open_existing(path: impl AsRef<Path>) -> Result<Ref<Self>, ()> {
31        let path_raw = path.as_ref().to_cstr();
32        let handle = NonNull::new(unsafe { BNOpenDatabase(path_raw.as_ptr()) }).ok_or(())?;
33        Ok(unsafe { Self::ref_from_raw(handle) })
34    }
35
36    /// Get a [`Snapshot`] by its `id`, or `None` if no snapshot with that `id` exists.
37    pub fn snapshot_by_id(&self, id: SnapshotId) -> Option<Ref<Snapshot>> {
38        let result = unsafe { BNGetDatabaseSnapshot(self.handle.as_ptr(), id.0) };
39        NonNull::new(result).map(|handle| unsafe { Snapshot::ref_from_raw(handle) })
40    }
41
42    /// Get a list of all snapshots in the database
43    pub fn snapshots(&self) -> Array<Snapshot> {
44        let mut count = 0;
45        let result = unsafe { BNGetDatabaseSnapshots(self.handle.as_ptr(), &mut count) };
46        assert!(!result.is_null());
47        unsafe { Array::new(result, count, ()) }
48    }
49
50    /// Get the current snapshot
51    pub fn current_snapshot(&self) -> Option<Ref<Snapshot>> {
52        let result = unsafe { BNGetDatabaseCurrentSnapshot(self.handle.as_ptr()) };
53        NonNull::new(result).map(|handle| unsafe { Snapshot::ref_from_raw(handle) })
54    }
55
56    /// Equivalent to [`Self::set_current_snapshot_id`].
57    pub fn set_current_snapshot(&self, value: &Snapshot) {
58        self.set_current_snapshot_id(value.id())
59    }
60
61    /// Sets the current snapshot to the [`SnapshotId`].
62    ///
63    /// **No** validation is done to ensure that the id is valid.
64    pub fn set_current_snapshot_id(&self, id: SnapshotId) {
65        unsafe { BNSetDatabaseCurrentSnapshot(self.handle.as_ptr(), id.0) }
66    }
67
68    pub fn write_snapshot_data(
69        &self,
70        parents: &[SnapshotId],
71        file: &BinaryView,
72        name: &str,
73        data: &KeyValueStore,
74        auto_save: bool,
75    ) -> SnapshotId {
76        self.write_snapshot_data_with_progress(
77            parents,
78            file,
79            name,
80            data,
81            auto_save,
82            NoProgressCallback,
83        )
84    }
85
86    pub fn write_snapshot_data_with_progress<P>(
87        &self,
88        parents: &[SnapshotId],
89        file: &BinaryView,
90        name: &str,
91        data: &KeyValueStore,
92        auto_save: bool,
93        mut progress: P,
94    ) -> SnapshotId
95    where
96        P: ProgressCallback,
97    {
98        let name_raw = name.to_cstr();
99        let name_ptr = name_raw.as_ptr();
100
101        let new_id = unsafe {
102            BNWriteDatabaseSnapshotData(
103                self.handle.as_ptr(),
104                // SAFETY: SnapshotId is just i64
105                parents.as_ptr() as *mut _,
106                parents.len(),
107                file.handle,
108                name_ptr,
109                data.handle.as_ptr(),
110                auto_save,
111                &mut progress as *mut P as *mut c_void,
112                Some(P::cb_progress_callback),
113            )
114        };
115
116        SnapshotId(new_id)
117    }
118
119    /// Trim a snapshot's contents in the database but leave the parent/child hierarchy intact.
120    ///
121    /// NOTE: Future references to this snapshot will return `false` for [`Database::snapshot_has_data`]
122    pub fn trim_snapshot(&self, id: SnapshotId) -> Result<(), ()> {
123        if unsafe { BNTrimDatabaseSnapshot(self.handle.as_ptr(), id.0) } {
124            Ok(())
125        } else {
126            Err(())
127        }
128    }
129
130    /// Remove a snapshot in the database by id, deleting its contents and references.
131    /// Attempting to remove a snapshot with children will raise an exception.
132    pub fn remove_snapshot(&self, id: SnapshotId) -> Result<(), ()> {
133        if unsafe { BNRemoveDatabaseSnapshot(self.handle.as_ptr(), id.0) } {
134            Ok(())
135        } else {
136            Err(())
137        }
138    }
139    pub fn has_global(&self, key: &str) -> bool {
140        let key_raw = key.to_cstr();
141        unsafe { BNDatabaseHasGlobal(self.handle.as_ptr(), key_raw.as_ptr()) != 0 }
142    }
143
144    /// Get a list of keys for all globals in the database
145    pub fn global_keys(&self) -> Array<BnString> {
146        let mut count = 0;
147        let result = unsafe { BNGetDatabaseGlobalKeys(self.handle.as_ptr(), &mut count) };
148        assert!(!result.is_null());
149        unsafe { Array::new(result, count, ()) }
150    }
151
152    /// Get a dictionary of all globals
153    pub fn globals(&self) -> HashMap<String, BnString> {
154        self.global_keys()
155            .iter()
156            .filter_map(|key| Some((key.to_string(), self.read_global(key)?)))
157            .collect()
158    }
159
160    /// Get a specific global by key
161    pub fn read_global(&self, key: &str) -> Option<BnString> {
162        let key_raw = key.to_cstr();
163        let result = unsafe { BNReadDatabaseGlobal(self.handle.as_ptr(), key_raw.as_ptr()) };
164        unsafe { NonNull::new(result).map(|_| BnString::from_raw(result)) }
165    }
166
167    /// Write a global into the database
168    pub fn write_global(&self, key: &str, value: &str) -> bool {
169        let key_raw = key.to_cstr();
170        let value_raw = value.to_cstr();
171        unsafe { BNWriteDatabaseGlobal(self.handle.as_ptr(), key_raw.as_ptr(), value_raw.as_ptr()) }
172    }
173
174    /// Get a specific global by key, as a binary buffer
175    pub fn read_global_data(&self, key: &str) -> Option<DataBuffer> {
176        let key_raw = key.to_cstr();
177        let result = unsafe { BNReadDatabaseGlobalData(self.handle.as_ptr(), key_raw.as_ptr()) };
178        NonNull::new(result).map(|_| DataBuffer::from_raw(result))
179    }
180
181    /// Write a binary buffer into a global in the database
182    pub fn write_global_data(&self, key: &str, value: &DataBuffer) -> bool {
183        let key_raw = key.to_cstr();
184        unsafe { BNWriteDatabaseGlobalData(self.handle.as_ptr(), key_raw.as_ptr(), value.as_raw()) }
185    }
186
187    /// Get the backing analysis cache kvs
188    pub fn analysis_cache(&self) -> Ref<KeyValueStore> {
189        let result = unsafe { BNReadDatabaseAnalysisCache(self.handle.as_ptr()) };
190        unsafe { KeyValueStore::ref_from_raw(NonNull::new(result).unwrap()) }
191    }
192
193    #[deprecated(note = "Use crate::file_metadata::FileMetadata::reopen_moved_database instead")]
194    /// Closes then reopens the database.
195    pub fn reload_connection(&self) {}
196
197    pub fn write_analysis_cache(&self, val: &KeyValueStore) -> Result<(), ()> {
198        if unsafe { BNWriteDatabaseAnalysisCache(self.handle.as_ptr(), val.handle.as_ptr()) } {
199            Ok(())
200        } else {
201            Err(())
202        }
203    }
204
205    pub fn snapshot_has_data(&self, id: SnapshotId) -> bool {
206        unsafe { BNSnapshotHasData(self.handle.as_ptr(), id.0) }
207    }
208}
209
210impl Debug for Database {
211    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
212        f.debug_struct("Database")
213            .field("current_snapshot", &self.current_snapshot())
214            .field("snapshot_count", &self.snapshots().len())
215            .field("globals", &self.globals())
216            .field("analysis_cache", &self.analysis_cache())
217            .finish()
218    }
219}
220
221impl ToOwned for Database {
222    type Owned = Ref<Self>;
223
224    fn to_owned(&self) -> Self::Owned {
225        unsafe { RefCountable::inc_ref(self) }
226    }
227}
228
229unsafe impl RefCountable for Database {
230    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
231        Ref::new(Self {
232            handle: NonNull::new(BNNewDatabaseReference(handle.handle.as_ptr())).unwrap(),
233        })
234    }
235
236    unsafe fn dec_ref(handle: &Self) {
237        BNFreeDatabase(handle.handle.as_ptr());
238    }
239}