tsafe-core 1.0.12

//! Team vault — age-encrypted shared secret store for multi-user environments.
//!
//! Team vaults use X25519 age encryption so each team member's public key can
//! decrypt the shared vault independently.  Membership and key rotation are
//! tracked in the vault's metadata.

use std::collections::HashMap;

use chrono::Utc;

use crate::age_crypto;
use crate::crypto::{self, KeyPurpose, KeySchedule, VaultKey};
use crate::errors::{SafeError, SafeResult};
use crate::rbac::RbacProfile;
use crate::vault::{KdfParams, SecretEntry, VaultChallenge, VaultFile, VAULT_CHALLENGE_PLAINTEXT};

const TEAM_SCHEMA: &str = "tsafe/vault/v2";
const TEAM_KEY_SCHEDULE: KeySchedule = KeySchedule::HkdfSha256V1;

/// Create a new team vault file encrypted to the given age recipient public keys.
/// Returns the VaultFile and the randomly generated DEK (for the caller to use).
pub fn create_team_vault(recipients: &[String]) -> SafeResult<(VaultFile, VaultKey)> {
    create_team_vault_with_access_profile(recipients, RbacProfile::ReadWrite)
}

/// Create a new team vault file with an explicit access profile.
pub fn create_team_vault_with_access_profile(
    recipients: &[String],
    access_profile: RbacProfile,
) -> SafeResult<(VaultFile, VaultKey)> {
    access_profile.ensure_write_allowed()?;
    if recipients.is_empty() {
        return Err(SafeError::Crypto {
            context: "at least one recipient is required".into(),
        });
    }
    let parsed = age_crypto::parse_recipients(recipients)?;

    // Generate a random 256-bit data encryption key (DEK).
    let dek_bytes = crypto::random_salt(); // 32 bytes
    let dek = VaultKey::from_bytes(dek_bytes);
    let cipher = crypto::default_vault_cipher();

    // Wrap DEK for all recipients using age.
    let wrapped = age_crypto::encrypt_to_recipients(&parsed, dek.as_bytes())?;

    // Create vault challenge encrypted with a purpose-scoped subkey.
    let (ch_nonce, ch_ct) = crypto::encrypt_with_key_schedule(
        &dek,
        TEAM_KEY_SCHEDULE,
        KeyPurpose::VaultChallenge,
        cipher,
        VAULT_CHALLENGE_PLAINTEXT,
    )?;

    let now = Utc::now();
    let file = VaultFile {
        schema: TEAM_SCHEMA.to_string(),
        kdf: KdfParams {
            algorithm: "age".to_string(),
            m_cost: 0,
            t_cost: 0,
            p_cost: 0,
            salt: String::new(),
        },
        cipher: cipher.as_str().to_string(),
        vault_challenge: VaultChallenge {
            nonce: crypto::encode_b64(&ch_nonce),
            ciphertext: crypto::encode_b64(&ch_ct),
        },
        created_at: now,
        updated_at: now,
        secrets: HashMap::new(),
        age_recipients: recipients.to_vec(),
        wrapped_dek: Some(crypto::encode_b64(&wrapped)),
    };

    Ok((file, dek))
}

/// Open a team vault by unwrapping the DEK with an age identity.
pub fn unwrap_dek(file: &VaultFile, identities: &[Box<dyn age::Identity>]) -> SafeResult<VaultKey> {
    let wrapped_b64 = file.wrapped_dek.as_ref().ok_or_else(|| SafeError::Crypto {
        context: "not a team/age vault — no wrapped_dek".into(),
    })?;
    let wrapped = crypto::decode_b64(wrapped_b64)?;
    let dek_bytes = age_crypto::decrypt_with_identities(identities, &wrapped)?;
    if dek_bytes.len() != 32 {
        return Err(SafeError::Crypto {
            context: format!("DEK has wrong length: expected 32, got {}", dek_bytes.len()),
        });
    }
    let mut arr = [0u8; 32];
    arr.copy_from_slice(&dek_bytes);
    Ok(VaultKey::from_bytes(arr))
}

/// Add a recipient to a team vault. Re-wraps the DEK for all recipients (including the new one).
pub fn add_member(
    file: &mut VaultFile,
    new_recipient: &str,
    identities: &[Box<dyn age::Identity>],
) -> SafeResult<()> {
    add_member_with_access_profile(file, new_recipient, identities, RbacProfile::ReadWrite)
}

/// Add a recipient to a team vault under an explicit access profile.
pub fn add_member_with_access_profile(
    file: &mut VaultFile,
    new_recipient: &str,
    identities: &[Box<dyn age::Identity>],
    access_profile: RbacProfile,
) -> SafeResult<()> {
    access_profile.ensure_write_allowed()?;
    let _dek = unwrap_dek(file, identities)?;

    if file.age_recipients.contains(&new_recipient.to_string()) {
        return Err(SafeError::Crypto {
            context: format!("recipient already exists: {new_recipient}"),
        });
    }
    file.age_recipients.push(new_recipient.to_string());

    // Re-wrap DEK for all recipients.
    let parsed = age_crypto::parse_recipients(&file.age_recipients)?;
    let wrapped = age_crypto::encrypt_to_recipients(&parsed, _dek.as_bytes())?;
    file.wrapped_dek = Some(crypto::encode_b64(&wrapped));
    file.updated_at = Utc::now();

    Ok(())
}

/// Remove a recipient and re-key the vault (generate new DEK, re-encrypt all secrets).
/// The removed member could have cached the old DEK, so all secrets must be re-encrypted.
pub fn remove_member(
    file: &mut VaultFile,
    remove_recipient: &str,
    identities: &[Box<dyn age::Identity>],
) -> SafeResult<()> {
    remove_member_with_access_profile(file, remove_recipient, identities, RbacProfile::ReadWrite)
}

/// Remove a recipient and re-key the vault under an explicit access profile.
pub fn remove_member_with_access_profile(
    file: &mut VaultFile,
    remove_recipient: &str,
    identities: &[Box<dyn age::Identity>],
    access_profile: RbacProfile,
) -> SafeResult<()> {
    access_profile.ensure_write_allowed()?;
    let old_dek = unwrap_dek(file, identities)?;
    let old_cipher = crypto::parse_cipher_kind(&file.cipher)?;
    let challenge_nonce = crypto::decode_b64(&file.vault_challenge.nonce)?;
    let challenge_ct = crypto::decode_b64(&file.vault_challenge.ciphertext)?;
    let old_schedule = crypto::detect_key_schedule(
        &old_dek,
        KeyPurpose::VaultChallenge,
        old_cipher,
        &challenge_nonce,
        &challenge_ct,
        VAULT_CHALLENGE_PLAINTEXT,
    )?;

    file.age_recipients.retain(|r| r != remove_recipient);
    if file.age_recipients.is_empty() {
        return Err(SafeError::Crypto {
            context: "cannot remove the last recipient".into(),
        });
    }

    // Generate new DEK.
    let new_dek_bytes = crypto::random_salt();
    let new_dek = VaultKey::from_bytes(new_dek_bytes);
    let new_cipher = crypto::default_vault_cipher();

    // Re-encrypt all secrets under the new DEK.
    let mut new_secrets = HashMap::with_capacity(file.secrets.len());
    for (key, entry) in &file.secrets {
        // Decrypt with old DEK.
        let nonce = crypto::decode_b64(&entry.nonce)?;
        let ct = crypto::decode_b64(&entry.ciphertext)?;
        let pt = crypto::decrypt_with_key_schedule(
            &old_dek,
            old_schedule,
            KeyPurpose::SecretData,
            old_cipher,
            &nonce,
            &ct,
        )?;

        // Re-encrypt with the current HKDF-scoped schedule.
        let (new_nonce, new_ct) = crypto::encrypt_with_key_schedule(
            &new_dek,
            TEAM_KEY_SCHEDULE,
            KeyPurpose::SecretData,
            new_cipher,
            &pt,
        )?;

        // Re-encrypt history entries too.
        let mut new_history = Vec::new();
        for h in &entry.history {
            let hn = crypto::decode_b64(&h.nonce)?;
            let hct = crypto::decode_b64(&h.ciphertext)?;
            let hpt = crypto::decrypt_with_key_schedule(
                &old_dek,
                old_schedule,
                KeyPurpose::SecretData,
                old_cipher,
                &hn,
                &hct,
            )?;
            let (nhn, nhct) = crypto::encrypt_with_key_schedule(
                &new_dek,
                TEAM_KEY_SCHEDULE,
                KeyPurpose::SecretData,
                new_cipher,
                &hpt,
            )?;
            new_history.push(crate::vault::HistoryEntry {
                nonce: crypto::encode_b64(&nhn),
                ciphertext: crypto::encode_b64(&nhct),
                updated_at: h.updated_at,
            });
        }

        new_secrets.insert(
            key.clone(),
            SecretEntry {
                nonce: crypto::encode_b64(&new_nonce),
                ciphertext: crypto::encode_b64(&new_ct),
                created_at: entry.created_at,
                updated_at: entry.updated_at,
                tags: entry.tags.clone(),
                history: new_history,
            },
        );
    }
    file.secrets = new_secrets;

    // Re-encrypt vault challenge with the HKDF-scoped schedule.
    let (ch_nonce, ch_ct) = crypto::encrypt_with_key_schedule(
        &new_dek,
        TEAM_KEY_SCHEDULE,
        KeyPurpose::VaultChallenge,
        new_cipher,
        VAULT_CHALLENGE_PLAINTEXT,
    )?;
    file.vault_challenge = VaultChallenge {
        nonce: crypto::encode_b64(&ch_nonce),
        ciphertext: crypto::encode_b64(&ch_ct),
    };
    file.cipher = new_cipher.as_str().to_string();

    // Wrap new DEK for remaining recipients.
    let parsed = age_crypto::parse_recipients(&file.age_recipients)?;
    let wrapped = age_crypto::encrypt_to_recipients(&parsed, new_dek.as_bytes())?;
    file.wrapped_dek = Some(crypto::encode_b64(&wrapped));
    file.updated_at = Utc::now();

    Ok(())
}

/// List the current team members (recipient public keys).
pub fn members(file: &VaultFile) -> &[String] {
    &file.age_recipients
}

/// Check if a vault is a team vault (has age recipients).
pub fn is_team_vault(file: &VaultFile) -> bool {
    !file.age_recipients.is_empty() && file.wrapped_dek.is_some()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::age_crypto;
    use crate::crypto::CipherKind;
    use crate::vault::HistoryEntry;

    fn identities_from(secret: &str) -> Vec<Box<dyn age::Identity>> {
        age::IdentityFile::from_buffer(secret.as_bytes())
            .unwrap()
            .into_identities()
            .unwrap()
    }

    #[test]
    fn create_team_vault_uses_hkdf_scoped_challenge() {
        let (secret, recipient) = age_crypto::generate_identity();
        let identities = identities_from(&secret);
        let (file, _dek) = create_team_vault(&[recipient]).unwrap();
        let dek = unwrap_dek(&file, &identities).unwrap();
        let challenge_nonce = crypto::decode_b64(&file.vault_challenge.nonce).unwrap();
        let challenge_ct = crypto::decode_b64(&file.vault_challenge.ciphertext).unwrap();

        assert!(matches!(
            crypto::decrypt_for_cipher(
                crypto::default_vault_cipher(),
                &dek,
                &challenge_nonce,
                &challenge_ct
            ),
            Err(SafeError::DecryptionFailed)
        ));
        assert_eq!(
            crypto::decrypt_with_key_schedule(
                &dek,
                KeySchedule::HkdfSha256V1,
                KeyPurpose::VaultChallenge,
                crypto::default_vault_cipher(),
                &challenge_nonce,
                &challenge_ct
            )
            .unwrap(),
            VAULT_CHALLENGE_PLAINTEXT
        );
    }

    #[test]
    fn remove_member_migrates_legacy_team_vault_to_hkdf_schedule() {
        let (secret1, recipient1) = age_crypto::generate_identity();
        let (_secret2, recipient2) = age_crypto::generate_identity();
        let identities = identities_from(&secret1);
        let recipients = vec![recipient1.clone(), recipient2.clone()];
        let parsed_recipients = age_crypto::parse_recipients(&recipients).unwrap();

        let dek = VaultKey::from_bytes(crypto::random_salt());
        let wrapped =
            age_crypto::encrypt_to_recipients(&parsed_recipients, dek.as_bytes()).unwrap();
        let now = Utc::now();
        let (challenge_nonce, challenge_ct) =
            crypto::encrypt(&dek, VAULT_CHALLENGE_PLAINTEXT).unwrap();
        let (secret_nonce, secret_ct) = crypto::encrypt(&dek, b"legacy-team-secret").unwrap();
        let (history_nonce, history_ct) = crypto::encrypt(&dek, b"legacy-history").unwrap();
        let mut secrets = HashMap::new();
        secrets.insert(
            "TEAM_SECRET".into(),
            SecretEntry {
                nonce: crypto::encode_b64(&secret_nonce),
                ciphertext: crypto::encode_b64(&secret_ct),
                created_at: now,
                updated_at: now,
                tags: HashMap::new(),
                history: vec![HistoryEntry {
                    nonce: crypto::encode_b64(&history_nonce),
                    ciphertext: crypto::encode_b64(&history_ct),
                    updated_at: now,
                }],
            },
        );

        let mut file = VaultFile {
            schema: TEAM_SCHEMA.to_string(),
            kdf: KdfParams {
                algorithm: "age".to_string(),
                m_cost: 0,
                t_cost: 0,
                p_cost: 0,
                salt: String::new(),
            },
            cipher: CipherKind::XChaCha20Poly1305.as_str().to_string(),
            vault_challenge: VaultChallenge {
                nonce: crypto::encode_b64(&challenge_nonce),
                ciphertext: crypto::encode_b64(&challenge_ct),
            },
            created_at: now,
            updated_at: now,
            secrets,
            age_recipients: recipients,
            wrapped_dek: Some(crypto::encode_b64(&wrapped)),
        };

        remove_member(&mut file, &recipient2, &identities).unwrap();

        let new_dek = unwrap_dek(&file, &identities).unwrap();
        let new_challenge_nonce = crypto::decode_b64(&file.vault_challenge.nonce).unwrap();
        let new_challenge_ct = crypto::decode_b64(&file.vault_challenge.ciphertext).unwrap();
        let new_cipher = crypto::parse_cipher_kind(&file.cipher).unwrap();
        assert!(matches!(
            crypto::decrypt_for_cipher(
                new_cipher,
                &new_dek,
                &new_challenge_nonce,
                &new_challenge_ct
            ),
            Err(SafeError::DecryptionFailed)
        ));
        assert_eq!(
            crypto::decrypt_with_key_schedule(
                &new_dek,
                KeySchedule::HkdfSha256V1,
                KeyPurpose::VaultChallenge,
                new_cipher,
                &new_challenge_nonce,
                &new_challenge_ct
            )
            .unwrap(),
            VAULT_CHALLENGE_PLAINTEXT
        );

        let entry = &file.secrets["TEAM_SECRET"];
        let nonce = crypto::decode_b64(&entry.nonce).unwrap();
        let ciphertext = crypto::decode_b64(&entry.ciphertext).unwrap();
        assert_eq!(
            crypto::decrypt_with_key_schedule(
                &new_dek,
                KeySchedule::HkdfSha256V1,
                KeyPurpose::SecretData,
                new_cipher,
                &nonce,
                &ciphertext
            )
            .unwrap(),
            b"legacy-team-secret"
        );
        let history = &entry.history[0];
        let history_nonce = crypto::decode_b64(&history.nonce).unwrap();
        let history_ct = crypto::decode_b64(&history.ciphertext).unwrap();
        assert_eq!(
            crypto::decrypt_with_key_schedule(
                &new_dek,
                KeySchedule::HkdfSha256V1,
                KeyPurpose::SecretData,
                new_cipher,
                &history_nonce,
                &history_ct
            )
            .unwrap(),
            b"legacy-history"
        );
    }

    #[cfg(feature = "fips")]
    #[test]
    fn fips_build_creates_aes256gcm_team_vaults() {
        let (_secret, recipient) = age_crypto::generate_identity();
        let (file, _dek) = create_team_vault(&[recipient]).unwrap();
        assert_eq!(file.cipher, CipherKind::Aes256Gcm.as_str());
    }

    // ── add_member ────────────────────────────────────────────────────────────

    #[test]
    fn add_member_allows_new_member_to_unwrap_dek() {
        let (secret1, recipient1) = age_crypto::generate_identity();
        let (secret2, recipient2) = age_crypto::generate_identity();
        let identities1 = identities_from(&secret1);
        let identities2 = identities_from(&secret2);

        let (mut file, _) = create_team_vault(&[recipient1]).unwrap();
        add_member(&mut file, &recipient2, &identities1).unwrap();

        // Both members can now unwrap the DEK.
        assert!(unwrap_dek(&file, &identities1).is_ok());
        assert!(unwrap_dek(&file, &identities2).is_ok());
        assert_eq!(file.age_recipients.len(), 2);
    }

    #[test]
    fn add_member_rejects_duplicate_recipient() {
        let (secret, recipient) = age_crypto::generate_identity();
        let identities = identities_from(&secret);
        let (mut file, _) = create_team_vault(std::slice::from_ref(&recipient)).unwrap();

        let result = add_member(&mut file, &recipient, &identities);
        assert!(matches!(result, Err(SafeError::Crypto { .. })));
    }

    // ── members ───────────────────────────────────────────────────────────────

    #[test]
    fn members_returns_current_recipient_list() {
        let (_secret, recipient) = age_crypto::generate_identity();
        let (file, _) = create_team_vault(std::slice::from_ref(&recipient)).unwrap();

        let m = members(&file);
        assert_eq!(m.len(), 1);
        assert_eq!(m[0], recipient);
    }

    #[test]
    fn members_reflects_add_member() {
        let (secret1, recipient1) = age_crypto::generate_identity();
        let (_secret2, recipient2) = age_crypto::generate_identity();
        let identities1 = identities_from(&secret1);

        let (mut file, _) = create_team_vault(std::slice::from_ref(&recipient1)).unwrap();
        add_member(&mut file, &recipient2, &identities1).unwrap();

        let m = members(&file);
        assert_eq!(m.len(), 2);
        assert!(m.contains(&recipient1));
        assert!(m.contains(&recipient2));
    }

    #[test]
    fn read_only_profile_rejects_team_mutations() {
        let (secret1, recipient1) = age_crypto::generate_identity();
        let (_secret2, recipient2) = age_crypto::generate_identity();
        let identities1 = identities_from(&secret1);
        let (mut file, _) = create_team_vault(std::slice::from_ref(&recipient1)).unwrap();

        assert!(matches!(
            add_member_with_access_profile(
                &mut file,
                &recipient2,
                &identities1,
                RbacProfile::ReadOnly
            ),
            Err(SafeError::InvalidVault { .. })
        ));
        assert!(matches!(
            remove_member_with_access_profile(
                &mut file,
                &recipient1,
                &identities1,
                RbacProfile::ReadOnly
            ),
            Err(SafeError::InvalidVault { .. })
        ));
        assert!(matches!(
            create_team_vault_with_access_profile(&[recipient1], RbacProfile::ReadOnly),
            Err(SafeError::InvalidVault { .. })
        ));
    }

    // ── is_team_vault ─────────────────────────────────────────────────────────

    #[test]
    fn is_team_vault_returns_true_for_team_vault() {
        let (_secret, recipient) = age_crypto::generate_identity();
        let (file, _) = create_team_vault(&[recipient]).unwrap();
        assert!(is_team_vault(&file));
    }

    #[test]
    fn is_team_vault_returns_false_when_no_recipients() {
        let file = VaultFile {
            schema: "tsafe/vault/v1".into(),
            kdf: KdfParams {
                algorithm: "argon2id".into(),
                m_cost: 65536,
                t_cost: 3,
                p_cost: 4,
                salt: String::new(),
            },
            cipher: "xchacha20poly1305".into(),
            vault_challenge: VaultChallenge {
                nonce: String::new(),
                ciphertext: String::new(),
            },
            created_at: Utc::now(),
            updated_at: Utc::now(),
            secrets: HashMap::new(),
            age_recipients: vec![],
            wrapped_dek: None,
        };
        assert!(!is_team_vault(&file));
    }

    // ── error paths ───────────────────────────────────────────────────────────

    #[test]
    fn create_team_vault_with_no_recipients_errors() {
        let result = create_team_vault(&[]);
        assert!(matches!(result, Err(SafeError::Crypto { .. })));
    }

    #[test]
    fn remove_last_member_errors() {
        let (secret, recipient) = age_crypto::generate_identity();
        let identities = identities_from(&secret);
        let (mut file, _) = create_team_vault(std::slice::from_ref(&recipient)).unwrap();

        let result = remove_member(&mut file, &recipient, &identities);
        assert!(matches!(result, Err(SafeError::Crypto { .. })));
    }
}