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
9pub struct SqliteDB {
11 conn: Connection,
12}
13
14impl SqliteDB {
15 pub fn new(conn: Connection) -> SqliteDB {
16 SqliteDB { conn }
17 }
18 pub fn empty() -> SqliteDB {
20 SqliteDB { conn: sqlite::open(":memory:").unwrap() }
21 }
22 pub fn virtual_file(mut file: Box<dyn VirtualFile>) -> ForensicResult<SqliteDB> {
25 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}