tryaudex_core/

keystore.rs

use aes_gcm::{
    aead::{Aead, KeyInit},
    Aes256Gcm, Nonce,
};
use serde::{Deserialize, Serialize};

use crate::error::{AvError, Result};

const KEYRING_SERVICE: &str = "audex";
const KEYRING_USER: &str = "credential-encryption-key";
const NONCE_LEN: usize = 12;

static CACHED_KEY: std::sync::OnceLock<[u8; 32]> = std::sync::OnceLock::new();

/// Encrypted blob stored on disk instead of plaintext JSON.
#[derive(Debug, Serialize, Deserialize)]
pub struct EncryptedBlob {
    /// Base64-encoded nonce
    pub nonce: String,
    /// Base64-encoded ciphertext
    pub ciphertext: String,
}

/// Get or create the encryption key from the OS keychain.
/// Falls back to a machine-derived key if keychain is unavailable.
/// Key is cached per-process to ensure consistency.
pub(crate) fn get_or_create_key() -> [u8; 32] {
    *CACHED_KEY.get_or_init(|| {
        // Try OS keychain first
        if let Ok(key) = load_key_from_keychain() {
            return key;
        }

        // Try to generate and store a new key
        if let Ok(key) = generate_and_store_key() {
            return key;
        }

        // Fallback: derive from machine-specific data
        derive_fallback_key()
    })
}

fn load_key_from_keychain() -> std::result::Result<[u8; 32], ()> {
    let entry = keyring::Entry::new(KEYRING_SERVICE, KEYRING_USER).map_err(|_| ())?;
    let secret = entry.get_password().map_err(|_| ())?;
    use base64::Engine;
    let bytes = base64::engine::general_purpose::STANDARD
        .decode(&secret)
        .map_err(|_| ())?;
    if bytes.len() != 32 {
        return Err(());
    }
    let mut key = [0u8; 32];
    key.copy_from_slice(&bytes);
    Ok(key)
}

fn generate_and_store_key() -> std::result::Result<[u8; 32], ()> {
    use rand::Rng;
    let mut key = [0u8; 32];
    rand::rng().fill(&mut key);

    use base64::Engine;
    let encoded = base64::engine::general_purpose::STANDARD.encode(key);

    let entry = keyring::Entry::new(KEYRING_SERVICE, KEYRING_USER).map_err(|_| ())?;
    entry.set_password(&encoded).map_err(|_| ())?;

    Ok(key)
}

/// Derive a key from machine-specific data plus a per-user random salt.
///
/// The salt is generated once and stored in `~/.config/audex/keystore-salt`
/// with owner-only permissions (0600). Without the salt file, the key cannot
/// be reproduced even though the other inputs (USER, home_dir, machine-id)
/// are publicly readable.
///
/// This is a fallback for environments without a keychain daemon (headless CI,
/// containers). For production use, ensure an OS keychain is available or set
/// AUDEX_HMAC_KEY for audit integrity.
///
/// # Panics
///
/// Panics (via `expect`) when the salt file cannot be created, because
/// proceeding with only public inputs would allow any local user to recompute
/// the key. Set `AUDEX_HMAC_KEY` to bypass the keystore in environments with
/// a read-only filesystem.
fn derive_fallback_key() -> [u8; 32] {
    use sha2::Digest;

    eprintln!(
        "  \x1b[33m⚠\x1b[0m OS keychain unavailable — using fallback encryption key. \
         Strength depends on ~/.config/audex/keystore-salt being private."
    );

    let salt = load_or_create_salt().unwrap_or_else(|| {
        // The salt is the sole source of entropy that prevents local users
        // from recomputing the fallback key. Without it the key is derived
        // entirely from public inputs (USER, home dir, machine-id) and
        // offers no real protection. Fail hard rather than silently produce
        // a weak key.
        //
        // If you are running in a container or on a read-only filesystem,
        // set the AUDEX_HMAC_KEY environment variable to supply an
        // explicit high-entropy key instead.
        panic!(
            "audex: keystore salt file could not be created \
             (~/.config/audex/keystore-salt). Proceeding without it would \
             derive the encryption key from public inputs only, which any \
             local user could recompute. Set the AUDEX_HMAC_KEY environment \
             variable to provide an explicit key, or ensure the audex config \
             directory is writable."
        );
    });

    let mut hasher = sha2::Sha256::new();
    hasher.update(b"audex-fallback-key-v2");

    // Mix in the per-user random salt — the primary entropy source.
    hasher.update(&salt);

    // Mix in username
    if let Ok(user) = std::env::var("USER").or_else(|_| std::env::var("LOGNAME")) {
        hasher.update(user.as_bytes());
    }

    // Mix in home directory
    if let Some(home) = dirs::home_dir() {
        hasher.update(home.to_string_lossy().as_bytes());
    }

    // Mix in machine-id if available (Linux)
    if let Ok(machine_id) = std::fs::read_to_string("/etc/machine-id") {
        hasher.update(machine_id.trim().as_bytes());
    }

    hasher.finalize().into()
}

/// Load or create a 32-byte random salt in `~/.config/audex/keystore-salt`.
fn load_or_create_salt() -> Option<Vec<u8>> {
    let dir = dirs::config_dir()?.join("audex");
    let path = dir.join("keystore-salt");

    // Try to read existing salt
    if let Ok(data) = std::fs::read(&path) {
        if data.len() == 32 {
            return Some(data);
        }
    }

    // Generate new salt
    use rand::Rng;
    let mut salt = [0u8; 32];
    rand::rng().fill(&mut salt);

    // Store with restrictive permissions
    let _ = std::fs::create_dir_all(&dir);
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let _ = std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700));
    }
    #[cfg(unix)]
    let write_ok = {
        use std::io::Write;
        use std::os::unix::fs::OpenOptionsExt;
        match std::fs::OpenOptions::new()
            .write(true)
            .create_new(true)
            .mode(0o600)
            .open(&path)
        {
            Ok(mut file) => file.write_all(&salt).is_ok(),
            Err(_) => false,
        }
    };
    #[cfg(not(unix))]
    let write_ok = std::fs::write(&path, salt).is_ok();

    if write_ok {
        Some(salt.to_vec())
    } else {
        tracing::warn!("Failed to write keystore salt file; fallback key has reduced entropy");
        None
    }
}

/// Encrypt data using AES-256-GCM with a key from the OS keychain.
pub fn encrypt(plaintext: &[u8]) -> Result<EncryptedBlob> {
    let key = get_or_create_key();
    let cipher = Aes256Gcm::new_from_slice(&key)
        .map_err(|e| AvError::InvalidPolicy(format!("Encryption init failed: {}", e)))?;

    use rand::Rng;
    let mut nonce_bytes = [0u8; NONCE_LEN];
    rand::rng().fill(&mut nonce_bytes);
    let nonce = Nonce::from_slice(&nonce_bytes);

    let ciphertext = cipher
        .encrypt(nonce, plaintext)
        .map_err(|e| AvError::InvalidPolicy(format!("Encryption failed: {}", e)))?;

    use base64::Engine;
    Ok(EncryptedBlob {
        nonce: base64::engine::general_purpose::STANDARD.encode(nonce_bytes),
        ciphertext: base64::engine::general_purpose::STANDARD.encode(ciphertext),
    })
}

/// Decrypt data using AES-256-GCM with a key from the OS keychain.
pub fn decrypt(blob: &EncryptedBlob) -> Result<Vec<u8>> {
    let key = get_or_create_key();
    let cipher = Aes256Gcm::new_from_slice(&key)
        .map_err(|e| AvError::InvalidPolicy(format!("Decryption init failed: {}", e)))?;

    use base64::Engine;
    let nonce_bytes = base64::engine::general_purpose::STANDARD
        .decode(&blob.nonce)
        .map_err(|e| AvError::InvalidPolicy(format!("Invalid nonce: {}", e)))?;
    let ciphertext = base64::engine::general_purpose::STANDARD
        .decode(&blob.ciphertext)
        .map_err(|e| AvError::InvalidPolicy(format!("Invalid ciphertext: {}", e)))?;

    if nonce_bytes.len() != NONCE_LEN {
        return Err(AvError::InvalidPolicy("Invalid nonce length".to_string()));
    }

    let nonce = Nonce::from_slice(&nonce_bytes);
    cipher
        .decrypt(nonce, ciphertext.as_ref())
        .map_err(|e| AvError::InvalidPolicy(format!("Decryption failed: {}", e)))
}

/// Encrypt a serializable value and write to a file.
///
/// R6-M52: write to a temporary file in the same directory, then
/// atomically rename into place. A crash or power loss during the
/// write previously left the target file truncated (zero bytes or
/// partial JSON), which `decrypt_from_file` would fail to parse —
/// silently destroying the cached SSO session or credential. The
/// rename is atomic on POSIX and near-atomic on Windows/NTFS.
pub fn encrypt_to_file<T: Serialize>(path: &std::path::Path, value: &T) -> Result<()> {
    let json = serde_json::to_vec(value)?;
    let blob = encrypt(&json)?;
    let blob_json = serde_json::to_string(&blob)?;

    // Build a temp path in the same directory so rename doesn't cross
    // filesystem boundaries.
    let parent = path.parent().unwrap_or_else(|| std::path::Path::new("."));
    let file_name = path
        .file_name()
        .unwrap_or_else(|| std::ffi::OsStr::new("tmp"));
    let tmp_path = parent.join(format!(".{}.tmp", file_name.to_string_lossy()));

    // Write with restrictive permissions (owner-only) since these files
    // contain encrypted credentials and SSO sessions.
    #[cfg(unix)]
    {
        use std::os::unix::fs::OpenOptionsExt;
        let mut file = std::fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .mode(0o600)
            .open(&tmp_path)?;
        std::io::Write::write_all(&mut file, blob_json.as_bytes())?;
        // Flush to disk before the rename so the data is durable.
        use std::io::Write;
        file.flush()?;
    }
    #[cfg(not(unix))]
    {
        std::fs::write(&tmp_path, &blob_json)?;
    }

    std::fs::rename(&tmp_path, path).map_err(|e| {
        // Clean up the temp file on rename failure.
        let _ = std::fs::remove_file(&tmp_path);
        AvError::InvalidPolicy(format!(
            "Failed to atomically replace {}: {}",
            path.display(),
            e
        ))
    })?;

    Ok(())
}

/// Read and decrypt a value from a file.
///
/// R6-M51: cap the read at 1 MiB. The previous unbounded
/// `read_to_string` would follow a symlink to /dev/urandom or a
/// multi-GB log and OOM the process. Encrypted blobs are a few KB;
/// 1 MiB is generous.
pub fn decrypt_from_file<T: for<'de> Deserialize<'de>>(
    path: &std::path::Path,
) -> Result<Option<T>> {
    if !path.exists() {
        return Ok(None);
    }

    const MAX_ENCRYPTED_BLOB_BYTES: u64 = 1024 * 1024;
    let mut f = std::fs::File::open(path)?;
    let mut blob_json = String::new();
    use std::io::Read;
    f.by_ref()
        .take(MAX_ENCRYPTED_BLOB_BYTES + 1)
        .read_to_string(&mut blob_json)?;
    if blob_json.len() as u64 > MAX_ENCRYPTED_BLOB_BYTES {
        return Err(AvError::InvalidPolicy(format!(
            "Encrypted blob at {} exceeds {} byte limit",
            path.display(),
            MAX_ENCRYPTED_BLOB_BYTES
        )));
    }

    // Try to decrypt as encrypted blob
    if let Ok(blob) = serde_json::from_str::<EncryptedBlob>(&blob_json) {
        if !blob.nonce.is_empty() && !blob.ciphertext.is_empty() {
            let plaintext = decrypt(&blob)?;
            let value: T = serde_json::from_slice(&plaintext)?;
            return Ok(Some(value));
        }
    }

    // Fallback: try reading as plaintext JSON (migration from unencrypted)
    match serde_json::from_str::<T>(&blob_json) {
        Ok(value) => Ok(Some(value)),
        Err(e) => Err(AvError::InvalidPolicy(format!(
            "Failed to read cached data: {}",
            e
        ))),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encrypt_decrypt_roundtrip() {
        let plaintext = b"secret credential data here";
        let blob = encrypt(plaintext).unwrap();

        // Verify it's not plaintext
        assert_ne!(blob.ciphertext.as_bytes(), plaintext);

        let decrypted = decrypt(&blob).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_encrypt_decrypt_json_value() {
        use serde_json::json;
        let value = json!({
            "access_key_id": "AKIAIOSFODNN7EXAMPLE",
            "secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
            "session_token": "FwoGZXIvYXdzEBY..."
        });

        let json_bytes = serde_json::to_vec(&value).unwrap();
        let blob = encrypt(&json_bytes).unwrap();
        let decrypted = decrypt(&blob).unwrap();
        let parsed: serde_json::Value = serde_json::from_slice(&decrypted).unwrap();

        assert_eq!(parsed["access_key_id"], "AKIAIOSFODNN7EXAMPLE");
    }

    #[test]
    fn test_different_encryptions_differ() {
        let plaintext = b"same data";
        let blob1 = encrypt(plaintext).unwrap();
        let blob2 = encrypt(plaintext).unwrap();
        // Different nonces should produce different ciphertexts
        assert_ne!(blob1.ciphertext, blob2.ciphertext);
        assert_ne!(blob1.nonce, blob2.nonce);
    }

    #[test]
    fn test_tampered_ciphertext_fails() {
        let blob = encrypt(b"secret").unwrap();
        let mut tampered = blob;
        // Flip a character in the ciphertext
        let mut chars: Vec<char> = tampered.ciphertext.chars().collect();
        if let Some(c) = chars.get_mut(5) {
            *c = if *c == 'A' { 'B' } else { 'A' };
        }
        tampered.ciphertext = chars.into_iter().collect();
        assert!(decrypt(&tampered).is_err());
    }

    #[test]
    fn test_encrypted_blob_serialization() {
        let blob = encrypt(b"test data").unwrap();
        let json = serde_json::to_string(&blob).unwrap();
        let parsed: EncryptedBlob = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.nonce, blob.nonce);
        assert_eq!(parsed.ciphertext, blob.ciphertext);

        let decrypted = decrypt(&parsed).unwrap();
        assert_eq!(decrypted, b"test data");
    }

    #[test]
    fn test_fallback_key_deterministic() {
        let key1 = derive_fallback_key();
        let key2 = derive_fallback_key();
        assert_eq!(key1, key2);
    }
}