featherChat/warzone/crates/warzone-client/src/storage.rs

//! Local sled database: sessions, pre-keys, message history.

use anyhow::{Context, Result};
use warzone_protocol::ratchet::RatchetState;
use warzone_protocol::types::Fingerprint;
use x25519_dalek::StaticSecret;

pub struct LocalDb {
    sessions: sled::Tree,
    pre_keys: sled::Tree,
    contacts: sled::Tree,
    history: sled::Tree,
    sender_keys: sled::Tree,
    _db: sled::Db,
}

impl LocalDb {
    pub fn open() -> Result<Self> {
        let path = crate::keystore::data_dir().join("db");
        let db = match sled::open(&path) {
            Ok(db) => db,
            Err(e) => {
                let err_str = e.to_string();
                if err_str.contains("WouldBlock") || err_str.contains("lock") {
                    eprintln!("Error: Database is locked by another warzone process.");
                    eprintln!("  DB path: {}", path.display());
                    eprintln!();
                    eprintln!("  Check for running processes:");
                    eprintln!("    ps aux | grep warzone-client");
                    eprintln!();
                    eprintln!("  To force unlock (if no other process is running):");
                    eprintln!("    rm -rf {}", path.display());
                    eprintln!("    (This deletes sessions — you'll need to re-establish them)");
                    anyhow::bail!("database locked by another process");
                }
                return Err(e).context("failed to open local database");
            }
        };
        let sessions = db.open_tree("sessions")?;
        let pre_keys = db.open_tree("pre_keys")?;
        let contacts = db.open_tree("contacts")?;
        let history = db.open_tree("history")?;
        let sender_keys = db.open_tree("sender_keys")?;
        Ok(LocalDb {
            sessions,
            pre_keys,
            contacts,
            history,
            sender_keys,
            _db: db,
        })
    }

    /// Save a ratchet session for a peer.
    pub fn save_session(&self, peer: &Fingerprint, state: &RatchetState) -> Result<()> {
        let key = peer.to_hex();
        let data = state.serialize_versioned()
            .map_err(|e| anyhow::anyhow!("{}", e))?;
        self.sessions.insert(key.as_bytes(), data)?;
        self.sessions.flush()?;
        Ok(())
    }

    /// Delete a ratchet session for a peer (used for session recovery).
    pub fn delete_session(&self, peer: &Fingerprint) -> Result<()> {
        let key = peer.to_hex();
        self.sessions.remove(key.as_bytes())?;
        self.sessions.flush()?;
        Ok(())
    }

    /// Load a ratchet session for a peer.
    pub fn load_session(&self, peer: &Fingerprint) -> Result<Option<RatchetState>> {
        let key = peer.to_hex();
        match self.sessions.get(key.as_bytes())? {
            Some(data) => {
                let state = RatchetState::deserialize_versioned(&data)
                    .map_err(|e| anyhow::anyhow!("{}", e))?;
                Ok(Some(state))
            }
            None => Ok(None),
        }
    }

    /// Store the signed pre-key secret (for X3DH respond).
    pub fn save_signed_pre_key(&self, id: u32, secret: &StaticSecret) -> Result<()> {
        let key = format!("spk:{}", id);
        self.pre_keys
            .insert(key.as_bytes(), secret.to_bytes().as_slice())?;
        self.pre_keys.flush()?;
        Ok(())
    }

    /// Load the signed pre-key secret.
    pub fn load_signed_pre_key(&self, id: u32) -> Result<Option<StaticSecret>> {
        let key = format!("spk:{}", id);
        match self.pre_keys.get(key.as_bytes())? {
            Some(data) => {
                let mut bytes = [0u8; 32];
                bytes.copy_from_slice(&data);
                Ok(Some(StaticSecret::from(bytes)))
            }
            None => Ok(None),
        }
    }

    /// Store a one-time pre-key secret.
    pub fn save_one_time_pre_key(&self, id: u32, secret: &StaticSecret) -> Result<()> {
        let key = format!("otpk:{}", id);
        self.pre_keys
            .insert(key.as_bytes(), secret.to_bytes().as_slice())?;
        self.pre_keys.flush()?;
        Ok(())
    }

    /// Return the next available OTPK ID (one past the highest stored).
    pub fn next_otpk_id(&self) -> u32 {
        let mut max_id: Option<u32> = None;
        for item in self.pre_keys.iter() {
            if let Ok((k, _)) = item {
                let key_str = String::from_utf8_lossy(&k);
                if let Some(id_str) = key_str.strip_prefix("otpk:") {
                    if let Ok(id) = id_str.parse::<u32>() {
                        max_id = Some(max_id.map_or(id, |m: u32| m.max(id)));
                    }
                }
            }
        }
        max_id.map_or(0, |m| m + 1)
    }

    /// Load and remove a one-time pre-key secret.
    pub fn take_one_time_pre_key(&self, id: u32) -> Result<Option<StaticSecret>> {
        let key = format!("otpk:{}", id);
        match self.pre_keys.remove(key.as_bytes())? {
            Some(data) => {
                let mut bytes = [0u8; 32];
                bytes.copy_from_slice(&data);
                self.pre_keys.flush()?;
                Ok(Some(StaticSecret::from(bytes)))
            }
            None => Ok(None),
        }
    }

    // ── Sender Keys ──

    /// Save a sender key for a (sender, group) pair.
    pub fn save_sender_key(
        &self,
        sender_fp: &str,
        group_name: &str,
        key: &warzone_protocol::sender_keys::SenderKey,
    ) -> Result<()> {
        let db_key = format!("sk:{}:{}", sender_fp, group_name);
        let data = bincode::serialize(key).context("failed to serialize sender key")?;
        self.sender_keys.insert(db_key.as_bytes(), data)?;
        self.sender_keys.flush()?;
        Ok(())
    }

    /// Load a sender key for a (sender, group) pair.
    pub fn load_sender_key(
        &self,
        sender_fp: &str,
        group_name: &str,
    ) -> Result<Option<warzone_protocol::sender_keys::SenderKey>> {
        let db_key = format!("sk:{}:{}", sender_fp, group_name);
        match self.sender_keys.get(db_key.as_bytes())? {
            Some(data) => {
                let key = bincode::deserialize(&data)
                    .context("failed to deserialize sender key")?;
                Ok(Some(key))
            }
            None => Ok(None),
        }
    }

    // ── Contacts ──

    /// Add or update a contact. Called on send/receive.
    pub fn touch_contact(&self, fingerprint: &str, alias: Option<&str>) -> Result<()> {
        let fp = fingerprint.chars().filter(|c| c.is_ascii_hexdigit()).collect::<String>().to_lowercase();
        let now = chrono::Utc::now().timestamp();

        let mut record = match self.contacts.get(fp.as_bytes())? {
            Some(data) => serde_json::from_slice::<serde_json::Value>(&data).unwrap_or_default(),
            None => serde_json::json!({}),
        };
        let obj = record.as_object_mut().unwrap();
        obj.insert("fingerprint".into(), serde_json::json!(fp));
        obj.insert("last_seen".into(), serde_json::json!(now));
        if let Some(a) = alias {
            obj.insert("alias".into(), serde_json::json!(a));
        }
        if !obj.contains_key("first_seen") {
            obj.insert("first_seen".into(), serde_json::json!(now));
        }
        let count = obj.get("message_count").and_then(|v| v.as_u64()).unwrap_or(0);
        obj.insert("message_count".into(), serde_json::json!(count + 1));

        self.contacts.insert(fp.as_bytes(), serde_json::to_vec(&record)?)?;
        Ok(())
    }

    /// Get all contacts sorted by last_seen (most recent first).
    pub fn list_contacts(&self) -> Result<Vec<serde_json::Value>> {
        let mut contacts: Vec<serde_json::Value> = self.contacts.iter()
            .filter_map(|item| {
                item.ok().and_then(|(_, data)| serde_json::from_slice(&data).ok())
            })
            .collect();
        contacts.sort_by(|a, b| {
            let ta = a.get("last_seen").and_then(|v| v.as_i64()).unwrap_or(0);
            let tb = b.get("last_seen").and_then(|v| v.as_i64()).unwrap_or(0);
            tb.cmp(&ta)
        });
        Ok(contacts)
    }

    // ── Message History ──

    /// Store a message in local history.
    pub fn store_message(&self, peer_fp: &str, sender: &str, text: &str, is_self: bool) -> Result<()> {
        let fp = peer_fp.chars().filter(|c| c.is_ascii_hexdigit()).collect::<String>().to_lowercase();
        let now = chrono::Utc::now().timestamp();
        let id = uuid::Uuid::new_v4().to_string();

        let msg = serde_json::json!({
            "id": id,
            "peer": fp,
            "sender": sender,
            "text": text,
            "is_self": is_self,
            "timestamp": now,
        });

        // Key: hist:<peer_fp>:<timestamp>:<uuid> for ordered scan
        let key = format!("hist:{}:{}:{}", fp, now, id);
        self.history.insert(key.as_bytes(), serde_json::to_vec(&msg)?)?;
        Ok(())
    }

    /// Get message history with a peer (most recent N messages).
    pub fn get_history(&self, peer_fp: &str, limit: usize) -> Result<Vec<serde_json::Value>> {
        let fp = peer_fp.chars().filter(|c| c.is_ascii_hexdigit()).collect::<String>().to_lowercase();
        let prefix = format!("hist:{}:", fp);

        let mut messages: Vec<serde_json::Value> = self.history
            .scan_prefix(prefix.as_bytes())
            .filter_map(|item| {
                item.ok().and_then(|(_, data)| serde_json::from_slice(&data).ok())
            })
            .collect();

        // Take last N
        if messages.len() > limit {
            messages = messages.split_off(messages.len() - limit);
        }
        Ok(messages)
    }

    /// Export all data as JSON (for encrypted backup).
    pub fn export_all(&self) -> Result<serde_json::Value> {
        let mut sessions = serde_json::Map::new();
        for item in self.sessions.iter() {
            if let Ok((k, v)) = item {
                let key = String::from_utf8_lossy(&k).to_string();
                sessions.insert(key, serde_json::json!(base64::Engine::encode(
                    &base64::engine::general_purpose::STANDARD, &v
                )));
            }
        }

        let mut pre_keys = serde_json::Map::new();
        for item in self.pre_keys.iter() {
            if let Ok((k, v)) = item {
                let key = String::from_utf8_lossy(&k).to_string();
                pre_keys.insert(key, serde_json::json!(base64::Engine::encode(
                    &base64::engine::general_purpose::STANDARD, &v
                )));
            }
        }

        Ok(serde_json::json!({
            "version": 1,
            "sessions": sessions,
            "pre_keys": pre_keys,
        }))
    }

    /// Create an encrypted backup of all session data.
    /// Returns the backup file path.
    pub fn create_backup(&self, seed: &[u8; 32]) -> Result<std::path::PathBuf> {
        use std::io::Write;

        let backup_dir = crate::keystore::data_dir().join("backups");
        std::fs::create_dir_all(&backup_dir)?;

        // Collect all data
        let mut data = serde_json::Map::new();

        // Sessions
        let mut sessions = serde_json::Map::new();
        for item in self.sessions.iter() {
            if let Ok((key, value)) = item {
                let k = String::from_utf8_lossy(&key).to_string();
                sessions.insert(k, serde_json::Value::String(base64::Engine::encode(
                    &base64::engine::general_purpose::STANDARD, &value
                )));
            }
        }
        data.insert("sessions".into(), serde_json::Value::Object(sessions));

        // Contacts
        let mut contacts = serde_json::Map::new();
        for item in self.contacts.iter() {
            if let Ok((key, value)) = item {
                let k = String::from_utf8_lossy(&key).to_string();
                if let Ok(v) = serde_json::from_slice::<serde_json::Value>(&value) {
                    contacts.insert(k, v);
                }
            }
        }
        data.insert("contacts".into(), serde_json::Value::Object(contacts));

        // Sender keys
        let mut sender_keys = serde_json::Map::new();
        for item in self.sender_keys.iter() {
            if let Ok((key, value)) = item {
                let k = String::from_utf8_lossy(&key).to_string();
                sender_keys.insert(k, serde_json::Value::String(base64::Engine::encode(
                    &base64::engine::general_purpose::STANDARD, &value
                )));
            }
        }
        data.insert("sender_keys".into(), serde_json::Value::Object(sender_keys));

        // Serialize and encrypt
        let plaintext = serde_json::to_vec(&serde_json::Value::Object(data))?;
        let key_bytes = warzone_protocol::crypto::hkdf_derive(seed, b"", b"warzone-backup", 32);
        let mut key = [0u8; 32];
        key.copy_from_slice(&key_bytes);
        let encrypted = warzone_protocol::crypto::aead_encrypt(&key, &plaintext, b"warzone-backup-aad");

        // Write to temp file then rename (atomic)
        let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S").to_string();
        let filename = format!("backup_{}.wzbk", timestamp);
        let path = backup_dir.join(&filename);
        let tmp_path = backup_dir.join(format!(".{}.tmp", filename));

        let mut file = std::fs::File::create(&tmp_path)?;
        file.write_all(&encrypted)?;
        file.sync_all()?;
        std::fs::rename(&tmp_path, &path)?;

        // Rotate: keep last 3 backups
        let mut backups: Vec<_> = std::fs::read_dir(&backup_dir)?
            .filter_map(|e| e.ok())
            .filter(|e| e.file_name().to_string_lossy().ends_with(".wzbk"))
            .collect();
        backups.sort_by_key(|e| e.file_name());
        while backups.len() > 3 {
            if let Some(old) = backups.first() {
                let _ = std::fs::remove_file(old.path());
                backups.remove(0);
            }
        }

        Ok(path)
    }

    /// Import data from JSON backup (merges, doesn't overwrite existing).
    pub fn import_all(&self, data: &serde_json::Value) -> Result<usize> {
        let mut count = 0;

        if let Some(sessions) = data.get("sessions").and_then(|v| v.as_object()) {
            for (key, val) in sessions {
                if let Some(b64) = val.as_str() {
                    if let Ok(bytes) = base64::Engine::decode(
                        &base64::engine::general_purpose::STANDARD, b64
                    ) {
                        // Only import if not already present
                        if self.sessions.get(key.as_bytes())?.is_none() {
                            self.sessions.insert(key.as_bytes(), bytes)?;
                            count += 1;
                        }
                    }
                }
            }
        }

        if let Some(pre_keys) = data.get("pre_keys").and_then(|v| v.as_object()) {
            for (key, val) in pre_keys {
                if let Some(b64) = val.as_str() {
                    if let Ok(bytes) = base64::Engine::decode(
                        &base64::engine::general_purpose::STANDARD, b64
                    ) {
                        if self.pre_keys.get(key.as_bytes())?.is_none() {
                            self.pre_keys.insert(key.as_bytes(), bytes)?;
                            count += 1;
                        }
                    }
                }
            }
        }

        self.sessions.flush()?;
        self.pre_keys.flush()?;
        Ok(count)
    }
}