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::{
    decrypt_chunk, derive_kek, generate_dummy_slot, mnemonic_to_dek, validate_password,
    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() {
            if self.slots.len() > 2 {
                bail!(
                    "Ungültige Slot-Anzahl im Backup: {} (maximal 2 erlaubt)",
                    self.slots.len()
                );
            }
            let mut seen_ids = std::collections::HashSet::new();
            for s in &self.slots {
                if s.slot_id > 1 {
                    bail!("Ungültige Slot-ID {} im Backup: Es sind ausschließlich die Slot-IDs 0 und 1 erlaubt", s.slot_id);
                }
                if !seen_ids.insert(s.slot_id) {
                    bail!("Doppelte Slot-ID {} im Backup entdeckt", s.slot_id);
                }
                crate::crypto::validate_kdf_params(&s.kdf_params).map_err(|e| {
                    anyhow::anyhow!("KDF-Parameter in Slot {} ungültig: {e}", s.slot_id)
                })?;

                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")?;

                if s.slot_id == 0 {
                    if wrapped_dek.len() != 40 && wrapped_dek.len() != 32 {
                        bail!("Ungültige wrapped_dek-Länge in Slot 0: {} Bytes (erwartet: 40 oder 32)", wrapped_dek.len());
                    }
                } else if s.slot_id == 1 {
                    if wrapped_dek.len() != 72 && wrapped_dek.len() != 64 && wrapped_dek.len() != 32
                    {
                        bail!("Ungültige wrapped_dek-Länge in Slot 1: {} Bytes (erwartet: 72, 64 oder 32)", wrapped_dek.len());
                    }
                }

                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,
                });
            }

            if !seen_ids.contains(&0) {
                bail!("Ungültiges Backup: Slot 0 (Standard/Decoy Vault) fehlt");
            }
        } 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(())
}

/// Ermittelt automatisch, ob ein Notfallschlüssel zu Slot 0 (Standard/Decoy) oder Slot 1 (Hidden Vault) gehört (S-05).
pub fn detect_recovery_key_slot(db: &Database, dek: &[u8; 32], version: u32) -> Result<u32> {
    if let Ok(Some(carrier_id)) = db.find_carrier_node_id() {
        if let Ok(Some(chunk0)) = db.read_chunk(carrier_id, 0) {
            let is_dek0 = decrypt_chunk(
                dek,
                carrier_id,
                0,
                &chunk0.ciphertext,
                &chunk0.nonce,
                &chunk0.tag,
                version,
                chunk0.generation,
            )
            .is_ok();
            if is_dek0 {
                return Ok(0);
            } else {
                return Ok(1);
            }
        }
    }
    Ok(0)
}

/// 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,
    dek_0: Option<&[u8; 32]>,
) -> Result<()> {
    if !container_path.exists() {
        bail!(
            "Containerdatei '{}' existiert nicht.",
            container_path.display()
        );
    }

    validate_password(new_password)?;

    // 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 der echte DEK_0 benötigt (S-05)
        let dek_0_bytes = dek_0.ok_or_else(|| {
            anyhow::anyhow!("Für die Wiederherstellung von Slot 1 (Hidden Vault) wird der DEK von Slot 0 (Standard-Vault) benötigt.")
        })?;
        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 Dummy-Slot-Längenparität
    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, None)
}

/// Führt ein vollständiges Rekeying eines Containers durch (K-03).
/// Generiert einen neuen kryptografischen DEK, verschlüsselt alle Chunks um,
/// verpackt den neuen DEK mit dem bestehenden Passwort und gibt das neue
/// 24-Wort Notfallblatt (BIP-39 Mnemonic) zurück.
/// Das alte Notfallblatt wird dadurch unwiderruflich ungültig.
pub fn rekey_container(container_path: &Path, password: &str) -> Result<String> {
    if !container_path.exists() {
        bail!(
            "Containerdatei '{}' existiert nicht.",
            container_path.display()
        );
    }

    let db = Database::open(container_path).context("Konnte Container-Datenbank nicht öffnen")?;
    let meta = db
        .read_meta()
        .context("Konnte Container-Header nicht lesen")?;
    let keys = meta.authenticate(password).ok_or_else(|| {
        anyhow::anyhow!("Ungültiges Master-Passwort! Authentifizierung fehlgeschlagen.")
    })?;

    let old_dek = keys.dek().clone();
    let slot_id = keys.slot_id();
    let version = keys.version();
    let carrier_dek = keys.carrier_dek();
    let carrier_node_id = keys.carrier_node_id();

    // 1. Neuen DEK generieren
    let new_dek = crate::crypto::generate_dek();

    // 2. Chunks des aktiven Vaults umverschlüsseln
    db.rekey_vault(slot_id, &old_dek, &new_dek, version)?;

    // 3. Neuen KEK ableiten und DEK neu verpacken
    let new_salt = crate::crypto::generate_salt();
    let new_params = KdfParams::default();
    let new_kek = derive_kek(password, &new_salt, &new_params)?;

    let (new_wrapped_dek, new_nonce, new_tag) = if slot_id == 1 {
        let c_dek = carrier_dek
            .ok_or_else(|| anyhow::anyhow!("DEK_0 (Träger-Schlüssel) für Slot 1 fehlt"))?;
        let c_nid = carrier_node_id.unwrap_or(0);
        wrap_slot1_payload(&new_kek, &new_dek, &c_dek, c_nid)?
    } else {
        let c_nid = carrier_node_id.unwrap_or(0);
        wrap_slot0_payload(&new_kek, &new_dek, c_nid)?
    };

    db.update_slot_keys(
        slot_id,
        &new_salt,
        &new_params,
        &new_wrapped_dek,
        &new_nonce,
        &new_tag,
    )?;

    // 4. Metadaten-MAC mit neuem DEK aktualisieren
    db.set_active_slot_and_dek(slot_id, new_dek.clone());
    db.update_metadata_mac()?;
    db.checkpoint()?;

    // 5. Neues BIP-39 Notfallblatt generieren
    let new_phrase = crate::crypto::dek_to_mnemonic(&new_dek)?;
    Ok(new_phrase)
}

#[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: crate::crypto::MIN_MEMORY_COST_KIB,
            time_cost: crate::crypto::MIN_TIME_COST,
            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: crate::crypto::MIN_MEMORY_COST_KIB,
            time_cost: crate::crypto::MIN_TIME_COST,
            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);
    }
}