chat4n6_sqlite_forensics/
db.rs

1use crate::btree::walk_table_btree;
2use crate::header::{is_sqlite_header, DbHeader};
3use crate::record::RecoveredRecord;
4use anyhow::{bail, Result};
5use chat4n6_plugin_api::EvidenceSource;
6use std::collections::HashMap;
7
8pub struct ForensicEngine<'a> {
9    data: &'a [u8],
10    header: DbHeader,
11}
12
13impl<'a> ForensicEngine<'a> {
14    pub fn new(data: &'a [u8], _timezone_offset: Option<i32>) -> Result<Self> {
15        if !is_sqlite_header(data) {
16            bail!("not a SQLite database");
17        }
18        let header = DbHeader::parse(data).ok_or_else(|| anyhow::anyhow!("invalid DB header"))?;
19        Ok(Self { data, header })
20    }
21
22    /// Layer 1: traverse B-tree and recover all live records.
23    pub fn recover_layer1(&self) -> Result<Vec<RecoveredRecord>> {
24        let mut records = Vec::new();
25
26        // Build table name → root page mapping from sqlite_master
27        let table_roots = self.read_sqlite_master()?;
28
29        for (table_name, root_page) in &table_roots {
30            self.traverse_btree(*root_page, table_name, &mut records);
31        }
32
33        Ok(records)
34    }
35
36    fn traverse_btree(&self, root_page: u32, table: &str, records: &mut Vec<RecoveredRecord>) {
37        walk_table_btree(
38            self.data,
39            self.header.page_size,
40            root_page,
41            table,
42            EvidenceSource::Live,
43            records,
44        );
45    }
46
47    /// Read sqlite_master (page 1) to get table name → root page mappings.
48    fn read_sqlite_master(&self) -> Result<HashMap<String, u32>> {
49        let mut tables = HashMap::new();
50        let mut temp_records = Vec::new();
51        // sqlite_master is always rooted at page 1
52        self.traverse_btree(1, "sqlite_master", &mut temp_records);
53
54        for record in temp_records {
55            // sqlite_master columns: type, name, tbl_name, rootpage, sql
56            if record.values.len() < 5 {
57                continue;
58            }
59            use crate::record::SqlValue;
60            let obj_type = match &record.values[0] {
61                SqlValue::Text(s) => s.as_str(),
62                _ => continue,
63            };
64            if obj_type != "table" {
65                continue;
66            }
67            // col 1 = name (the object's own name — not tbl_name at col 2)
68            let name = match &record.values[1] {
69                SqlValue::Text(s) => s.clone(),
70                _ => continue,
71            };
72            let root_page = match &record.values[3] {
73                SqlValue::Int(n) => *n as u32,
74                _ => continue,
75            };
76            if root_page > 0 {
77                tables.insert(name, root_page);
78            }
79        }
80
81        Ok(tables)
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    use crate::record::SqlValue;
89
90    fn create_test_db() -> Vec<u8> {
91        let conn = rusqlite::Connection::open_in_memory().unwrap();
92        conn.execute_batch(
93            "CREATE TABLE messages (id INTEGER PRIMARY KEY, text TEXT, ts INTEGER);
94             INSERT INTO messages VALUES (1, 'hello world', 1710000000000);
95             INSERT INTO messages VALUES (2, 'foo bar', 1710000001000);",
96        )
97        .unwrap();
98        let tmp = tempfile::NamedTempFile::new().unwrap();
99        conn.backup(rusqlite::DatabaseName::Main, tmp.path(), None)
100            .unwrap();
101        std::fs::read(tmp.path()).unwrap()
102    }
103
104    #[test]
105    fn test_layer1_reads_live_records() {
106        let db_bytes = create_test_db();
107        let engine = ForensicEngine::new(&db_bytes, None).unwrap();
108        let results = engine.recover_layer1().unwrap();
109        let msgs: Vec<_> = results.iter().filter(|r| r.table == "messages").collect();
110        assert_eq!(msgs.len(), 2);
111        assert!(msgs
112            .iter()
113            .any(|r| r.values.get(1) == Some(&SqlValue::Text("hello world".into()))));
114    }
115
116    #[test]
117    fn test_layer1_rejects_non_sqlite() {
118        assert!(ForensicEngine::new(b"not a database", None).is_err());
119    }
120}
chat4n6_sqlite_forensics/db.rs

chat4n6_sqlite_forensics/
db.rs
