sanctum/src/recovery.rs

use std::fs;
use std::path::Path;
use std::time::{SystemTime, UNIX_EPOCH};

use anyhow::{bail, Context, Result};
use serde::{Deserialize, Serialize};

use crate::crypto::{
    derive_kek, generate_dummy_slot, mnemonic_to_dek, wrap_slot0_payload,
    wrap_slot1_payload, KdfParams, FORMAT_VERSION,
};
use crate::storage::{ContainerMeta, Database, SlotMeta};

pub const HEADER_BACKUP_MAGIC: &str = "SANCTUM_HEADER_BACKUP";
pub const CURRENT_BACKUP_VERSION: u32 = 1;

/// Struktur für einen gesicherten Header-Slot.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SlotBackup {
    pub slot_id: u32,
    pub version: u32,
    pub kdf_salt_hex: String,
    pub kdf_params: KdfParams,
    pub wrapped_dek_hex: String,
    pub header_nonce_hex: String,
    pub header_tag_hex: String,
}

/// Struktur für exportierte Header-Backups (.sanctum.hdr) im JSON-Format.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HeaderBackup {
    pub magic: String,
    pub backup_version: u32,
    pub container_format_version: u32,
    pub kdf_salt_hex: String,
    pub kdf_params: KdfParams,
    pub wrapped_dek_hex: String,
    pub header_nonce_hex: String,
    pub header_tag_hex: String,
    #[serde(default)]
    pub slots: Vec<SlotBackup>,
    pub created_at: u64,
}

impl HeaderBackup {
    pub fn from_meta(meta: &ContainerMeta) -> Self {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_secs())
            .unwrap_or(0);

        let slots: Vec<SlotBackup> = meta
            .slots
            .iter()
            .map(|s| SlotBackup {
                slot_id: s.slot_id,
                version: s.version,
                kdf_salt_hex: hex::encode(s.kdf_salt),
                kdf_params: s.kdf_params.clone(),
                wrapped_dek_hex: hex::encode(&s.wrapped_dek),
                header_nonce_hex: hex::encode(s.header_nonce),
                header_tag_hex: hex::encode(s.header_tag),
            })
            .collect();

        Self {
            magic: HEADER_BACKUP_MAGIC.to_string(),
            backup_version: CURRENT_BACKUP_VERSION,
            container_format_version: meta.version,
            kdf_salt_hex: hex::encode(meta.kdf_salt),
            kdf_params: meta.kdf_params.clone(),
            wrapped_dek_hex: hex::encode(&meta.wrapped_dek),
            header_nonce_hex: hex::encode(meta.header_nonce),
            header_tag_hex: hex::encode(meta.header_tag),
            slots,
            created_at: now,
        }
    }

    pub fn to_meta(&self) -> Result<ContainerMeta> {
        if self.magic != HEADER_BACKUP_MAGIC {
            bail!("Ungültige Header-Backup-Datei: Falsches Magic-Präfix");
        }

        let mut slots = Vec::new();

        if !self.slots.is_empty() {
            for s in &self.slots {
                let salt_bytes = hex::decode(&s.kdf_salt_hex)
                    .context("Ungültige Hex-Kodierung für KDF-Salt")?;
                if salt_bytes.len() != 16 {
                    bail!("Ungültige Salt-Länge im Backup: erwartet 16 Bytes, erhalten {}", salt_bytes.len());
                }
                let mut kdf_salt = [0u8; 16];
                kdf_salt.copy_from_slice(&salt_bytes);

                let wrapped_dek = hex::decode(&s.wrapped_dek_hex)
                    .context("Ungültige Hex-Kodierung für wrapped_dek")?;

                let nonce_bytes = hex::decode(&s.header_nonce_hex)
                    .context("Ungültige Hex-Kodierung für Header-Nonce")?;
                if nonce_bytes.len() != 12 {
                    bail!("Ungültige Nonce-Länge im Backup: erwartet 12 Bytes, erhalten {}", nonce_bytes.len());
                }
                let mut header_nonce = [0u8; 12];
                header_nonce.copy_from_slice(&nonce_bytes);

                let tag_bytes = hex::decode(&s.header_tag_hex)
                    .context("Ungültige Hex-Kodierung für Header-Tag")?;
                if tag_bytes.len() != 16 {
                    bail!("Ungültige Tag-Länge im Backup: erwartet 16 Bytes, erhalten {}", tag_bytes.len());
                }
                let mut header_tag = [0u8; 16];
                header_tag.copy_from_slice(&tag_bytes);

                slots.push(SlotMeta {
                    slot_id: s.slot_id,
                    version: s.version,
                    kdf_salt,
                    kdf_params: s.kdf_params.clone(),
                    wrapped_dek,
                    header_nonce,
                    header_tag,
                });
            }
        } else {
            // Fallback für alte Backups ohne slots-Array
            let salt_bytes = hex::decode(&self.kdf_salt_hex)
                .context("Ungültige Hex-Kodierung für KDF-Salt")?;
            if salt_bytes.len() != 16 {
                bail!("Ungültige Salt-Länge im Backup: erwartet 16 Bytes, erhalten {}", salt_bytes.len());
            }
            let mut kdf_salt = [0u8; 16];
            kdf_salt.copy_from_slice(&salt_bytes);

            let wrapped_dek = hex::decode(&self.wrapped_dek_hex)
                .context("Ungültige Hex-Kodierung für wrapped_dek")?;

            let nonce_bytes = hex::decode(&self.header_nonce_hex)
                .context("Ungültige Hex-Kodierung für Header-Nonce")?;
            if nonce_bytes.len() != 12 {
                bail!("Ungültige Nonce-Länge im Backup: erwartet 12 Bytes, erhalten {}", nonce_bytes.len());
            }
            let mut header_nonce = [0u8; 12];
            header_nonce.copy_from_slice(&nonce_bytes);

            let tag_bytes = hex::decode(&self.header_tag_hex)
                .context("Ungültige Hex-Kodierung für Header-Tag")?;
            if tag_bytes.len() != 16 {
                bail!("Ungültige Tag-Länge im Backup: erwartet 16 Bytes, erhalten {}", tag_bytes.len());
            }
            let mut header_tag = [0u8; 16];
            header_tag.copy_from_slice(&tag_bytes);

            let slot0 = SlotMeta {
                slot_id: 0,
                version: self.container_format_version,
                kdf_salt,
                kdf_params: self.kdf_params.clone(),
                wrapped_dek: wrapped_dek.clone(),
                header_nonce,
                header_tag,
            };
            slots.push(slot0);
        }

        let slot0 = &slots[0];
        Ok(ContainerMeta {
            version: slot0.version,
            kdf_salt: slot0.kdf_salt,
            kdf_params: slot0.kdf_params.clone(),
            wrapped_dek: slot0.wrapped_dek.clone(),
            header_nonce: slot0.header_nonce,
            header_tag: slot0.header_tag,
            slots,
        })
    }
}

/// Sichert die Header-Metadaten eines Containers in eine externe Backup-Datei (.sanctum.hdr).
pub fn export_header_backup(container_path: &Path, backup_path: &Path) -> Result<()> {
    if !container_path.exists() {
        bail!("Containerdatei '{}' existiert nicht.", container_path.display());
    }

    let db = Database::open(container_path)
        .context("Konnte Container-Datenbank zum Lesen des Headers nicht öffnen")?;
    let meta = db.read_meta().context("Konnte Container-Header nicht lesen")?;

    let backup = HeaderBackup::from_meta(&meta);
    let json_data = serde_json::to_string_pretty(&backup)
        .context("Fehler beim Serialisieren des Header-Backups")?;

    fs::write(backup_path, json_data)
        .with_context(|| format!("Konnte Backup-Datei '{}' nicht schreiben", backup_path.display()))?;

    Ok(())
}

/// Stellt den Container-Header aus einer Sicherungsdatei (.sanctum.hdr) wieder her.
pub fn restore_header_backup(container_path: &Path, backup_path: &Path) -> Result<()> {
    if !backup_path.exists() {
        bail!("Backup-Datei '{}' existiert nicht.", backup_path.display());
    }

    let content = fs::read_to_string(backup_path)
        .with_context(|| format!("Konnte Backup-Datei '{}' nicht lesen", backup_path.display()))?;

    let backup: HeaderBackup = serde_json::from_str(&content)
        .context("Ungültiges Backup-Dateiformat (JSON-Parsing fehlgeschlagen)")?;

    let meta = backup.to_meta()?;

    let db = Database::open(container_path)
        .context("Konnte Ziel-Containerdatei nicht öffnen")?;

    db.restore_meta(&meta)
        .context("Fehler beim Wiederherstellen der Header-Tabelle in der Datenbank")?;

    db.checkpoint().context("Fehler beim WAL-Checkpoint nach Header-Wiederherstellung")?;

    Ok(())
}

/// Rekonstruiert einen Slot des Container-Headers mithilfe des 24-Wort BIP-39 Notfallschlüssels
/// und initialisiert ein neues Master-Passwort für den entsprechenden Slot (0 = Decoy, 1 = Hidden Vault).
pub fn restore_slot_from_recovery_key(
    container_path: &Path,
    recovery_key: &str,
    new_password: &str,
    target_slot_id: u32,
) -> Result<()> {
    if !container_path.exists() {
        bail!("Containerdatei '{}' existiert nicht.", container_path.display());
    }

    if new_password.trim().is_empty() {
        bail!("Das neue Master-Passwort darf nicht leer sein.");
    }

    // 1. DEK aus 24-Wort-Phrase dekodieren & validieren
    let dek = mnemonic_to_dek(recovery_key)
        .context("Ungültiger 24-Wort Notfallschlüssel")?;

    // 2. Neuen KEK mit frischem Salt ableiten
    let mut salt = [0u8; 16];
    rand::RngCore::fill_bytes(&mut rand::rngs::OsRng, &mut salt);

    let kdf_params = KdfParams::default();
    let kek = derive_kek(new_password, &salt, &kdf_params)
        .context("Schlüsselableitung für neues Passwort fehlgeschlagen")?;

    let db = Database::open(container_path)
        .context("Konnte Container-Datenbank nicht öffnen")?;

    let carrier_node_id = db.find_carrier_node_id()?.unwrap_or(0);

    let mut existing_slots = db.read_slots().unwrap_or_default();

    let new_slot = if target_slot_id == 1 {
        // Für Slot 1 wird DEK_0 benötigt (z. B. aus vorhandenem Slot 0)
        let dek_0_bytes = [0u8; 32];
        let (wrapped, nonce, tag) = wrap_slot1_payload(&kek, &dek, &dek_0_bytes, carrier_node_id)?;
        SlotMeta {
            slot_id: 1,
            version: FORMAT_VERSION,
            kdf_salt: salt,
            kdf_params: kdf_params.clone(),
            wrapped_dek: wrapped,
            header_nonce: nonce,
            header_tag: tag,
        }
    } else {
        // Slot 0 (Standard / Decoy Vault) mit 40 Bytes für Modell A
        let (wrapped, nonce, tag) = wrap_slot0_payload(&kek, &dek, carrier_node_id)?;
        SlotMeta {
            slot_id: 0,
            version: FORMAT_VERSION,
            kdf_salt: salt,
            kdf_params: kdf_params.clone(),
            wrapped_dek: wrapped,
            header_nonce: nonce,
            header_tag: tag,
        }
    };

    // Slot ersetzen bzw. einfügen
    existing_slots.retain(|s| s.slot_id != target_slot_id);
    existing_slots.push(new_slot);
    existing_slots.sort_by_key(|s| s.slot_id);

    // Falls Slot 1 fehlt, Dummy-Slot 1 ergänzen für Plausible Deniability
    if !existing_slots.iter().any(|s| s.slot_id == 1) {
        let (dummy_dek, dummy_nonce, dummy_tag, dummy_salt) = generate_dummy_slot();
        let dummy_slot = SlotMeta {
            slot_id: 1,
            version: FORMAT_VERSION,
            kdf_salt: dummy_salt,
            kdf_params: KdfParams::default(),
            wrapped_dek: dummy_dek,
            header_nonce: dummy_nonce,
            header_tag: dummy_tag,
        };
        existing_slots.push(dummy_slot);
    }

    let slot0_meta = existing_slots
        .iter()
        .find(|s| s.slot_id == 0)
        .ok_or_else(|| anyhow::anyhow!("Slot 0 fehlt im Header"))?;

    let meta = ContainerMeta {
        version: FORMAT_VERSION,
        kdf_salt: slot0_meta.kdf_salt,
        kdf_params: slot0_meta.kdf_params.clone(),
        wrapped_dek: slot0_meta.wrapped_dek.clone(),
        header_nonce: slot0_meta.header_nonce,
        header_tag: slot0_meta.header_tag,
        slots: existing_slots,
    };

    db.restore_meta(&meta)
        .context("Fehler beim Schreiben des rekonstruierten Headers")?;

    db.checkpoint().context("Fehler beim WAL-Checkpoint nach Header-Rekonstruktion")?;

    Ok(())
}

/// Rekonstruiert den Decoy-Slot (Slot 0) mithilfe des 24-Wort BIP-39 Notfallschlüssels.
pub fn restore_header_from_recovery_key(
    container_path: &Path,
    recovery_key: &str,
    new_password: &str,
) -> Result<()> {
    restore_slot_from_recovery_key(container_path, recovery_key, new_password, 0)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::crypto::{generate_dek, generate_salt, unwrap_dek, wrap_dek};
    use std::path::PathBuf;

    #[test]
    fn test_header_backup_and_restore() {
        let temp_dir = std::env::temp_dir();
        let container_path: PathBuf = temp_dir.join(format!("test_backup_{}.sanctum", std::process::id()));
        let backup_path: PathBuf = temp_dir.join(format!("test_backup_{}.sanctum.hdr", std::process::id()));

        if container_path.exists() {
            let _ = fs::remove_file(&container_path);
        }
        if backup_path.exists() {
            let _ = fs::remove_file(&backup_path);
        }

        let password = "SuperSecretPassword2026!";
        let salt = generate_salt();
        let kdf_params = KdfParams {
            memory_cost: 1024,
            time_cost: 1,
            parallelism: 1,
        };
        let kek = derive_kek(password, &salt, &kdf_params).unwrap();
        let dek = generate_dek();
        let (wrapped_dek, nonce, tag) = wrap_dek(&kek, &dek).unwrap();

        let db = Database::open(&container_path).unwrap();
        db.init_schema(&salt, &kdf_params, &wrapped_dek, &nonce, &tag).unwrap();
        db.checkpoint().unwrap();

        // 1. Export
        export_header_backup(&container_path, &backup_path).expect("Export header");
        assert!(backup_path.exists());

        // 2. Header in DB gezielt zerstören/löschen
        let conn = rusqlite::Connection::open(&container_path).unwrap();
        conn.execute("DELETE FROM meta", []).unwrap();
        drop(conn);

        // Prüfen, dass Container jetzt unlesbar ist
        let broken_db = Database::open(&container_path).unwrap();
        assert!(broken_db.read_meta().is_err());
        drop(broken_db);

        // 3. Restore
        restore_header_backup(&container_path, &backup_path).expect("Restore header");

        // 4. Verifikation: Container wieder voll entschlüsselbar
        let restored_db = Database::open(&container_path).unwrap();
        let meta = restored_db.read_meta().expect("Read restored meta");
        let restored_kek = derive_kek(password, &meta.kdf_salt, &meta.kdf_params).unwrap();
        let active_dek = unwrap_dek(&restored_kek, &meta.wrapped_dek, &meta.header_nonce, &meta.header_tag)
            .expect("Unwrap restored DEK");
        assert_eq!(*dek, *active_dek);

        let _ = fs::remove_file(&container_path);
        let _ = fs::remove_file(&backup_path);
    }

    #[test]
    fn test_restore_from_recovery_key() {
        let temp_dir = std::env::temp_dir();
        let container_path: PathBuf = temp_dir.join(format!("test_rec_key_{}.sanctum", std::process::id()));

        if container_path.exists() {
            let _ = fs::remove_file(&container_path);
        }

        let old_password = "ForgottenOldPassword!";
        let new_password = "BrandNewRescuedPassword2026!";

        let salt = generate_salt();
        let kdf_params = KdfParams {
            memory_cost: 1024,
            time_cost: 1,
            parallelism: 1,
        };
        let kek = derive_kek(old_password, &salt, &kdf_params).unwrap();
        let dek = generate_dek();
        let (wrapped_dek, nonce, tag) = wrap_dek(&kek, &dek).unwrap();

        // 24-Wort Notfallschlüssel sichern
        let phrase = crate::crypto::dek_to_mnemonic(&dek).unwrap();

        let db = Database::open(&container_path).unwrap();
        db.init_schema(&salt, &kdf_params, &wrapped_dek, &nonce, &tag).unwrap();
        db.checkpoint().unwrap();

        // Header zerstören
        let conn = rusqlite::Connection::open(&container_path).unwrap();
        conn.execute("DELETE FROM meta", []).unwrap();
        drop(conn);

        // Mit 24-Wort Schlüssel und NEUEM Passwort wiederherstellen
        restore_header_from_recovery_key(&container_path, &phrase, new_password)
            .expect("Restore from recovery key");

        // Verifikation: Neues Passwort entschlüsselt originalen DEK
        let rescued_db = Database::open(&container_path).unwrap();
        let meta = rescued_db.read_meta().expect("Read rescued meta");
        let new_kek = derive_kek(new_password, &meta.kdf_salt, &meta.kdf_params).unwrap();
        let unwrapped = unwrap_dek(&new_kek, &meta.wrapped_dek, &meta.header_nonce, &meta.header_tag)
            .expect("Unwrap rescued DEK");
        assert_eq!(*dek, *unwrapped);

        let _ = fs::remove_file(&container_path);
    }
}