frnsc_sqlite/
lib.rs

1use std::{io::Write, time::{UNIX_EPOCH, SystemTime, Duration}};
2
3use forensic_rs::{
4    prelude::{ForensicError, ForensicResult},
5    traits::{sql::{ColumnType, ColumnValue, SqlDb, SqlStatement}, vfs::VirtualFile},
6};
7use sqlite::{Connection, Statement, OpenFlags};
8
9/// SQLite DB that implements the forensic SqlDb trait
10pub struct SqliteDB {
11    conn: Connection,
12}
13
14impl SqliteDB {
15    pub fn new(conn: Connection) -> SqliteDB {
16        SqliteDB { conn }
17    }
18    /// Create an empty in-memmory DB
19    pub fn empty() -> SqliteDB {
20        SqliteDB { conn: sqlite::open(":memory:").unwrap() }
21    }
22    /// Create a SQLite DB from a virtual file in ReadOnly and Serialized mode. The implementation copies the entire SQLite into a temp folder and opens it.
23    /// The alternative is create a custom VFS for SQLite. https://www.sqlite.org/vfs.html
24    pub fn virtual_file(mut file: Box<dyn VirtualFile>) -> ForensicResult<SqliteDB> {
25        // We need to copy the full file from the virtual filesystem into a temp file in the machine
26        let mut buffer = vec![0; 4096];
27        let millis = match SystemTime::now()
28            .duration_since(UNIX_EPOCH) {
29                Ok(v) => v,
30                Err(_) => Duration::from_secs(1)
31            }.subsec_nanos();
32        let file_name =format!("forensic_sqlite.{}.db", millis);
33        let temp_path = std::env::temp_dir().join(file_name);
34        let mut tmp_file = std::fs::File::create(&temp_path)?;
35        loop {
36            let readed = file.read(&mut buffer)?;
37            if readed == 0 {
38                break;
39            }
40            tmp_file.write_all(&buffer[0..readed])?;
41        }
42        let connection = match sqlite::Connection::open_with_flags(&temp_path.to_string_lossy()[..], OpenFlags::new().set_read_only().set_full_mutex()) {
43            Ok(v) => v,
44            Err(e) => return Err(ForensicError::Other(e.to_string()))
45        };
46        Ok(SqliteDB::new(connection))
47    }
48}
49
50impl SqlDb for SqliteDB {
51    fn prepare<'a>(&'a self, statement: &'a str) -> ForensicResult<Box<dyn SqlStatement + 'a>> {
52        Ok(Box::new(SqliteStatement::new(&self.conn, statement)?))
53    }
54
55    fn from_file(&self, file: Box<dyn VirtualFile>) -> ForensicResult<Box<dyn SqlDb>> {
56        Ok(Box::new(Self::virtual_file(file)?))
57    }
58    fn list_tables(&self) -> ForensicResult<Vec<String>> {
59        let mut ret = Vec::with_capacity(32);
60        let mut sts = self.prepare(r#"SELECT 
61        name
62    FROM 
63        sqlite_schema
64    WHERE 
65        type ='table' AND 
66        name NOT LIKE 'sqlite_%';"#)?;
67        loop {
68            if !sts.next()? {
69                break;
70            }
71            let name : String = sts.read(0)?.try_into()?;
72            ret.push(name);
73        }
74        Ok(ret)
75    }
76}
77
78pub struct SqliteStatement<'conn> {
79    stmt: Statement<'conn>,
80}
81impl<'conn> SqliteStatement<'conn> {
82    pub fn new(conn: &'conn Connection, statement: &str) -> ForensicResult<SqliteStatement<'conn>> {
83        Ok(Self {
84            stmt: match conn.prepare(statement) {
85                Ok(st) => st,
86                Err(e) => return Err(ForensicError::Other(e.to_string())),
87            },
88        })
89    }
90}
91
92impl<'conn> SqlStatement for SqliteStatement<'conn> {
93    fn column_count(&self) -> usize {
94        self.stmt.column_count()
95    }
96
97    fn column_name(&self, i: usize) -> Option<&str> {
98        match self.stmt.column_name(i) {
99            Ok(v) => Some(v),
100            Err(_) => None,
101        }
102    }
103
104    fn column_names(&self) -> Vec<&str> {
105        self.stmt.column_names().iter().map(|v| &v[..]).collect()
106    }
107
108    fn column_type(&self, i: usize) -> ColumnType {
109        let column_type = match self.stmt.column_type(i) {
110            Ok(v) => v,
111            Err(_e) => return ColumnType::Null,
112        };
113        match column_type {
114            sqlite::Type::Binary => ColumnType::Binary,
115            sqlite::Type::Float => ColumnType::Float,
116            sqlite::Type::Integer => ColumnType::Integer,
117            sqlite::Type::String => ColumnType::String,
118            sqlite::Type::Null => ColumnType::Null,
119        }
120    }
121
122    fn next(&mut self) -> ForensicResult<bool> {
123        match self.stmt.next() {
124            Ok(v) => Ok(match v {
125                sqlite::State::Row => true,
126                sqlite::State::Done => false,
127            }),
128            Err(e) => Err(ForensicError::Other(e.to_string())),
129        }
130    }
131
132    fn read(&self, i: usize) -> ForensicResult<ColumnValue> {
133        let column_type = match self.stmt.column_type(i) {
134            Ok(v) => v,
135            Err(e) => return Err(ForensicError::Other(e.to_string())),
136        };
137        match column_type {
138            sqlite::Type::Binary => match self.stmt.read(i) {
139                Ok(v) => Ok(ColumnValue::Binary(v)),
140                Err(e) => Err(ForensicError::Other(e.to_string())),
141            },
142            sqlite::Type::Float => match self.stmt.read(i) {
143                Ok(v) => Ok(ColumnValue::Float(v)),
144                Err(e) => Err(ForensicError::Other(e.to_string())),
145            },
146            sqlite::Type::Integer => match self.stmt.read(i) {
147                Ok(v) => Ok(ColumnValue::Integer(v)),
148                Err(e) => Err(ForensicError::Other(e.to_string())),
149            },
150            sqlite::Type::String => match self.stmt.read(i) {
151                Ok(v) => Ok(ColumnValue::String(v)),
152                Err(e) => Err(ForensicError::Other(e.to_string())),
153            },
154            sqlite::Type::Null => Ok(ColumnValue::Null),
155        }
156    }
157}
158
159#[cfg(test)]
160mod test_db_implementation {
161    use super::*;
162
163    use forensic_rs::{traits::{sql::{SqlStatement, SqlDb}, vfs::VirtualFileSystem}, prelude::ForensicResult};
164    use sqlite::Connection;
165
166    use crate::SqliteDB;
167
168    fn initialize_mem_db() -> Connection {
169        let connection = sqlite::open(":memory:").unwrap();
170        prepare_db(connection)
171    }
172    fn initialize_file_db() -> Connection {
173        let millis = match SystemTime::now()
174            .duration_since(UNIX_EPOCH) {
175                Ok(v) => v,
176                Err(_) => Duration::from_secs(1)
177            }.subsec_nanos();
178        let file_name =format!("forensic_sqlite.{}.db", millis);
179        let temp_path = std::env::temp_dir().join(file_name);
180        let connection = sqlite::open(&temp_path).unwrap();
181        prepare_db(connection)
182    }
183
184    fn prepare_db(connection : Connection) -> Connection {
185        connection
186            .execute(
187                "
188            CREATE TABLE users (name TEXT, age INTEGER);
189            INSERT INTO users VALUES ('Alice', 42);
190            INSERT INTO users VALUES ('Bob', 69);
191            ",
192            )
193            .unwrap();
194        connection
195    }
196    fn prepare_wrapper(connection: Connection) -> SqliteDB {
197        SqliteDB::new(connection)
198    }
199
200    #[test]
201    fn sqlite_in_memory() {
202        let conn = initialize_mem_db();
203        let w_conn = prepare_wrapper(conn);
204        let mut statement = w_conn.prepare("SELECT name, age FROM users;").unwrap();
205        test_database_content(statement.as_mut()).expect("Should not return error");
206    }
207
208    #[test]
209    fn sqlite_from_machine_file() {
210        let conn = initialize_file_db();
211        let w_conn = prepare_wrapper(conn);
212        let mut statement = w_conn.prepare("SELECT name, age FROM users;").unwrap();
213        test_database_content(statement.as_mut()).expect("Should not return error");
214    }
215
216    #[test]
217    fn sqlite_from_virtual_file() {
218        let millis = match SystemTime::now()
219            .duration_since(UNIX_EPOCH) {
220                Ok(v) => v,
221                Err(_) => Duration::from_secs(1)
222            }.as_millis();
223        let file_name =format!("forensic_sqlite.{}.db", millis);
224        let temp_path = std::env::temp_dir().join(file_name);
225        let connection = sqlite::open(&temp_path).unwrap();
226        prepare_db(connection);
227
228        let mut fs = forensic_rs::core::fs::StdVirtualFS::new();
229        let file = fs.open(&temp_path).unwrap();
230        let w_conn = SqliteDB::virtual_file(file).unwrap();
231        let mut statement = w_conn.prepare("SELECT name, age FROM users;").unwrap();
232        test_database_content(statement.as_mut()).expect("Should not return error");
233    }
234
235    fn test_database_content<'a>(statement: &mut dyn SqlStatement) -> ForensicResult<()> {
236        assert!(statement.next()?);
237        let name: String = statement.read(0)?.try_into()?;
238        let age: usize = statement.read(1)?.try_into()?;
239        assert_eq!("Alice", name);
240        assert_eq!(42, age);
241        assert!(statement.next()?);
242        let name: String = statement.read(0)?.try_into()?;
243        let age: usize = statement.read(1)?.try_into()?;
244        assert_eq!("Bob", name);
245        assert_eq!(69, age);
246        assert!(!statement.next()?);
247        Ok(())
248    }
249}