roboticus_core/

keystore.rs

use std::collections::HashMap;
use std::io::Write;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

use aes_gcm::aead::{Aead, KeyInit, OsRng};
use aes_gcm::{Aes256Gcm, Nonce};
use argon2::Argon2;
use chrono::Utc;
use rand::RngCore;
use serde::{Deserialize, Serialize};
use serde_json::json;
use zeroize::Zeroizing;

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

const SALT_LEN: usize = 16;
const NONCE_LEN: usize = 12;

/// Acquires the mutex, recovering from poison if a prior holder panicked.
///
/// Rationale: the keystore is an in-memory cache backed by an encrypted file.
/// If a thread panics mid-update the in-memory state may be stale but not
/// corrupt -- the on-disk file remains the source of truth and will be
/// re-read on the next `refresh_locked` call. Panicking here would make
/// the entire keystore permanently unavailable, which is worse than serving
/// a potentially stale cache.
fn lock_or_recover<T>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
    m.lock().unwrap_or_else(|e| e.into_inner())
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct KeystoreData {
    entries: HashMap<String, String>,
}

/// Combined in-memory state behind a single mutex, eliminating lock ordering
/// concerns that existed with the previous two-mutex design.
struct KeystoreState {
    entries: Option<HashMap<String, Zeroizing<String>>>,
    passphrase: Option<Zeroizing<String>>,
    last_file_fingerprint: Option<(std::time::SystemTime, u64)>,
}

#[derive(Clone)]
pub struct Keystore {
    path: PathBuf,
    state: Arc<Mutex<KeystoreState>>,
}

impl Keystore {
    pub fn new(path: impl Into<PathBuf>) -> Self {
        Self {
            path: path.into(),
            state: Arc::new(Mutex::new(KeystoreState {
                entries: None,
                passphrase: None,
                last_file_fingerprint: None,
            })),
        }
    }

    pub fn default_path() -> PathBuf {
        crate::home_dir().join(".roboticus").join("keystore.enc")
    }

    pub fn unlock(&self, passphrase: &str) -> Result<()> {
        if !self.path.exists() {
            let mut st = lock_or_recover(&self.state);
            st.entries = Some(HashMap::new());
            st.passphrase = Some(Zeroizing::new(passphrase.to_string()));
            st.last_file_fingerprint = None;
            drop(st);
            self.save()?;
            self.append_audit_event(
                "initialize",
                None,
                json!({
                    "result": "ok",
                    "details": "created new keystore file"
                }),
            )?;
            return Ok(());
        }
        let zeroized_entries = self.decrypt_entries(passphrase)?;
        let mut st = lock_or_recover(&self.state);
        st.entries = Some(zeroized_entries);
        st.passphrase = Some(Zeroizing::new(passphrase.to_string()));
        st.last_file_fingerprint = self.current_file_fingerprint();
        Ok(())
    }

    /// Unlock with a machine-derived passphrase.
    ///
    /// Strategy:
    /// 1. Try the current machine-id based passphrase (stable random ID).
    /// 2. Fall back to legacy hostname-based passphrases for migration.
    /// 3. On legacy success, re-key the keystore to the new machine-id passphrase.
    ///
    /// **Security note:** This provides convenience-only protection. The passphrase
    /// file is readable by any local process. Use `unlock()` with a user-supplied
    /// passphrase for secrets requiring real confidentiality.
    pub fn unlock_machine(&self) -> Result<()> {
        // Try the current machine-id passphrase first.
        let primary = machine_passphrase();
        if self.unlock(&primary).is_ok() {
            return Ok(());
        }

        // Try all legacy hostname-based passphrases (pre-rebrand ironclad-*,
        // post-rebrand roboticus-*, env vars, gethostname syscall).
        for legacy in legacy_passphrases() {
            if legacy != primary && self.unlock(&legacy).is_ok() {
                tracing::info!("keystore unlocked with legacy passphrase; migrating to machine-id");
                // Auto-migrate: re-key to the stable machine-id passphrase so
                // future unlocks don't depend on hostname stability.
                if let Err(e) = self.rekey(&primary) {
                    tracing::warn!(error = %e, "failed to migrate keystore to machine-id passphrase");
                } else {
                    tracing::info!("keystore migrated to machine-id passphrase");
                }
                return Ok(());
            }
        }

        // All strategies failed — return the error from the primary attempt.
        self.unlock(&primary)
    }

    pub fn is_unlocked(&self) -> bool {
        lock_or_recover(&self.state).entries.is_some()
    }

    pub fn get(&self, key: &str) -> Option<String> {
        let mut st = lock_or_recover(&self.state);
        if let Err(e) = self.refresh_locked(&mut st) {
            tracing::warn!(error = %e, "keystore refresh failed, using cached entries");
        }
        st.entries
            .as_ref()
            .and_then(|m| m.get(key).map(|v| (**v).clone()))
    }

    pub fn set(&self, key: &str, value: &str) -> Result<()> {
        let previous = {
            let mut st = lock_or_recover(&self.state);
            let entries = st
                .entries
                .as_mut()
                .ok_or_else(|| RoboticusError::Keystore("keystore is locked".into()))?;
            entries.insert(key.to_string(), Zeroizing::new(value.to_string()))
        };
        let save_res = self.save();
        let rolled_back = save_res.is_err();
        if rolled_back {
            let mut st = lock_or_recover(&self.state);
            if let Some(entries) = st.entries.as_mut() {
                if let Some(prev) = previous {
                    entries.insert(key.to_string(), prev);
                } else {
                    entries.remove(key);
                }
            }
        }
        let audit_res = self.append_audit_event(
            "set",
            Some(key),
            json!({
                "result": if save_res.is_ok() { "ok" } else { "error" },
                "rolled_back": rolled_back
            }),
        );
        match (save_res, audit_res) {
            (Err(e), _) => Err(e),
            (Ok(()), Err(e)) => Err(e),
            (Ok(()), Ok(())) => Ok(()),
        }
    }

    pub fn remove(&self, key: &str) -> Result<bool> {
        let removed = {
            let mut st = lock_or_recover(&self.state);
            let entries = st
                .entries
                .as_mut()
                .ok_or_else(|| RoboticusError::Keystore("keystore is locked".into()))?;
            entries.remove(key)
        };
        let existed = removed.is_some();
        if existed {
            let save_res = self.save();
            let rolled_back = save_res.is_err();
            if rolled_back {
                let mut st = lock_or_recover(&self.state);
                if let Some(entries) = st.entries.as_mut()
                    && let Some(prev) = removed
                {
                    entries.insert(key.to_string(), prev);
                }
            }
            let audit_res = self.append_audit_event(
                "remove",
                Some(key),
                json!({
                    "result": if save_res.is_ok() { "ok" } else { "error" },
                    "rolled_back": rolled_back
                }),
            );
            match (save_res, audit_res) {
                (Err(e), _) => return Err(e),
                (Ok(()), Err(e)) => return Err(e),
                (Ok(()), Ok(())) => {}
            }
        }
        Ok(existed)
    }

    pub fn list_keys(&self) -> Vec<String> {
        let mut st = lock_or_recover(&self.state);
        if let Err(e) = self.refresh_locked(&mut st) {
            tracing::warn!(error = %e, "keystore refresh failed, using cached entries");
        }
        st.entries
            .as_ref()
            .map(|m| m.keys().cloned().collect())
            .unwrap_or_default()
    }

    pub fn import(&self, new_entries: HashMap<String, String>) -> Result<usize> {
        let count = new_entries.len();
        let snapshot = {
            let mut st = lock_or_recover(&self.state);
            let entries = st
                .entries
                .as_mut()
                .ok_or_else(|| RoboticusError::Keystore("keystore is locked".into()))?;
            let before = entries.clone();
            entries.extend(new_entries.into_iter().map(|(k, v)| (k, Zeroizing::new(v))));
            before
        };
        let save_res = self.save();
        let rolled_back = save_res.is_err();
        if rolled_back {
            let mut st = lock_or_recover(&self.state);
            st.entries = Some(snapshot);
        }
        let audit_res = self.append_audit_event(
            "import",
            None,
            json!({
                "result": if save_res.is_ok() { "ok" } else { "error" },
                "count": count,
                "rolled_back": rolled_back
            }),
        );
        match (save_res, audit_res) {
            (Err(e), _) => return Err(e),
            (Ok(()), Err(e)) => return Err(e),
            (Ok(()), Ok(())) => {}
        }
        Ok(count)
    }

    pub fn lock(&self) {
        let mut st = lock_or_recover(&self.state);
        st.entries = None;
        st.passphrase = None;
    }

    /// Re-encrypt with a new passphrase. Must already be unlocked.
    pub fn rekey(&self, new_passphrase: &str) -> Result<()> {
        if !self.is_unlocked() {
            return Err(RoboticusError::Keystore("keystore is locked".into()));
        }
        let old_passphrase = {
            let mut st = lock_or_recover(&self.state);
            let prev = st.passphrase.clone();
            st.passphrase = Some(Zeroizing::new(new_passphrase.to_string()));
            prev
        };
        let save_res = self.save();
        let rolled_back = save_res.is_err();
        if rolled_back {
            let mut st = lock_or_recover(&self.state);
            st.passphrase = old_passphrase;
        }
        let audit_res = self.append_audit_event(
            "rekey",
            None,
            json!({
                "result": if save_res.is_ok() { "ok" } else { "error" },
                "rolled_back": rolled_back
            }),
        );
        match (save_res, audit_res) {
            (Err(e), _) => Err(e),
            (Ok(()), Err(e)) => Err(e),
            (Ok(()), Ok(())) => Ok(()),
        }
    }

    fn audit_log_path(&self) -> PathBuf {
        self.path.with_extension("audit.log")
    }

    fn append_audit_event(
        &self,
        operation: &str,
        key: Option<&str>,
        metadata: serde_json::Value,
    ) -> Result<()> {
        let audit_path = self.audit_log_path();
        if let Some(parent) = audit_path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(&audit_path)?;
        #[cfg(unix)]
        if let Ok(meta) = file.metadata() {
            use std::os::unix::fs::PermissionsExt;
            if meta.permissions().mode() & 0o777 != 0o600
                && let Err(e) =
                    std::fs::set_permissions(&audit_path, std::fs::Permissions::from_mode(0o600))
            {
                tracing::warn!(error = %e, path = %audit_path.display(), "failed to set keystore audit log permissions");
            }
        }

        let redacted_key = key.map(redact_key_name);
        let record = json!({
            "timestamp": Utc::now().to_rfc3339(),
            "operation": operation,
            "key": redacted_key,
            "pid": std::process::id(),
            "process": std::env::args().next().unwrap_or_else(|| "unknown".to_string()),
            "keystore_path": self.path,
            "metadata": metadata
        });
        file.write_all(record.to_string().as_bytes())?;
        file.write_all(b"\n")?;
        file.flush()?;
        Ok(())
    }

    fn save(&self) -> Result<()> {
        let st = lock_or_recover(&self.state);
        let entries = st
            .entries
            .as_ref()
            .ok_or_else(|| RoboticusError::Keystore("keystore is locked".into()))?;

        let passphrase = st
            .passphrase
            .as_ref()
            .ok_or_else(|| RoboticusError::Keystore("no passphrase available".into()))?;

        let salt = fresh_salt();
        let key = derive_key(passphrase, &salt)?;

        let store = KeystoreData {
            entries: entries
                .iter()
                .map(|(k, v)| (k.clone(), (**v).clone()))
                .collect(),
        };
        let plaintext = serde_json::to_vec(&store)?;

        let cipher = Aes256Gcm::new_from_slice(key.as_ref())
            .map_err(|e| RoboticusError::Keystore(e.to_string()))?;

        let mut nonce_bytes = [0u8; NONCE_LEN];
        OsRng.fill_bytes(&mut nonce_bytes);
        let nonce = Nonce::from_slice(&nonce_bytes);

        let ciphertext = cipher
            .encrypt(nonce, plaintext.as_ref())
            .map_err(|e| RoboticusError::Keystore(format!("encryption failed: {e}")))?;

        let mut out = Vec::with_capacity(SALT_LEN + NONCE_LEN + ciphertext.len());
        out.extend_from_slice(&salt);
        out.extend_from_slice(&nonce_bytes);
        out.extend_from_slice(&ciphertext);

        // Drop lock before filesystem I/O to avoid holding it during blocking ops.
        drop(st);

        if let Some(parent) = self.path.parent() {
            std::fs::create_dir_all(parent)?;
        }

        let tmp = self.path.with_extension("tmp");
        std::fs::write(&tmp, &out)?;

        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            std::fs::set_permissions(&tmp, std::fs::Permissions::from_mode(0o600))?;
        }

        std::fs::rename(&tmp, &self.path)?;
        let fingerprint = self.current_file_fingerprint();
        let mut st = lock_or_recover(&self.state);
        st.last_file_fingerprint = fingerprint;

        Ok(())
    }

    fn decrypt_entries(&self, passphrase: &str) -> Result<HashMap<String, Zeroizing<String>>> {
        let data = std::fs::read(&self.path)?;
        if data.len() < SALT_LEN + NONCE_LEN + 1 {
            return Err(RoboticusError::Keystore("corrupt keystore file".into()));
        }

        let salt = &data[..SALT_LEN];
        let nonce_bytes = &data[SALT_LEN..SALT_LEN + NONCE_LEN];
        let ciphertext = &data[SALT_LEN + NONCE_LEN..];

        let key = derive_key(passphrase, salt)?;
        let cipher = Aes256Gcm::new_from_slice(key.as_ref())
            .map_err(|e| RoboticusError::Keystore(e.to_string()))?;
        let nonce = Nonce::from_slice(nonce_bytes);

        let plaintext = cipher.decrypt(nonce, ciphertext).map_err(|_| {
            RoboticusError::Keystore("decryption failed (wrong passphrase?)".into())
        })?;

        let store: KeystoreData = serde_json::from_slice(&plaintext)
            .map_err(|e| RoboticusError::Keystore(format!("corrupt keystore data: {e}")))?;

        Ok(store
            .entries
            .into_iter()
            .map(|(k, v)| (k, Zeroizing::new(v)))
            .collect())
    }

    fn refresh_locked(&self, st: &mut KeystoreState) -> Result<()> {
        if st.entries.is_none() {
            return Ok(());
        }
        let Some(passphrase) = st.passphrase.as_ref() else {
            return Ok(());
        };
        if !self.path.exists() {
            return Ok(());
        }
        let current_fingerprint = self.current_file_fingerprint();
        if current_fingerprint.is_some() && st.last_file_fingerprint == current_fingerprint {
            return Ok(());
        }
        let refreshed = self.decrypt_entries(passphrase)?;
        st.entries = Some(refreshed);
        st.last_file_fingerprint = current_fingerprint;
        Ok(())
    }

    fn current_file_fingerprint(&self) -> Option<(std::time::SystemTime, u64)> {
        let meta = std::fs::metadata(&self.path).ok()?;
        let modified = meta.modified().ok()?;
        Some((modified, meta.len()))
    }
}

fn derive_key(passphrase: &str, salt: &[u8]) -> Result<Zeroizing<[u8; 32]>> {
    let params = argon2::Params::new(65536, 3, 1, Some(32))
        .map_err(|e| RoboticusError::Keystore(format!("argon2 params: {e}")))?;
    let argon2 = Argon2::new(argon2::Algorithm::Argon2id, argon2::Version::V0x13, params);

    let mut key = Zeroizing::new([0u8; 32]);
    argon2
        .hash_password_into(passphrase.as_bytes(), salt, key.as_mut())
        .map_err(|e| RoboticusError::Keystore(format!("key derivation failed: {e}")))?;
    Ok(key)
}

fn fresh_salt() -> [u8; SALT_LEN] {
    let mut salt = [0u8; SALT_LEN];
    OsRng.fill_bytes(&mut salt);
    salt
}

/// Redact a key name for audit logging: show the first 3 characters followed
/// by `***` so that logs are useful for debugging without exposing full names.
fn redact_key_name(key: &str) -> String {
    let visible: String = key.chars().take(3).collect();
    format!("{visible}***")
}

/// Returns the path to the persistent machine ID file.
///
/// In test mode (`ROBOTICUS_TEST_MACHINE_ID_DIR` set), uses a temp directory
/// to prevent tests from destroying the production keystore's encryption key.
fn machine_id_path() -> PathBuf {
    if let Ok(test_dir) = std::env::var("ROBOTICUS_TEST_MACHINE_ID_DIR") {
        return PathBuf::from(test_dir).join("machine-id");
    }
    crate::home_dir().join(".roboticus").join("machine-id")
}

/// Read or create a stable machine ID.
///
/// The machine ID is a random hex string generated on first use and persisted
/// at `~/.roboticus/machine-id`. Unlike hostname-based derivation, this value
/// never changes across network contexts, DHCP leases, machine renames, or
/// shell environments.
///
/// **Security note:** This provides convenience-only protection (same as the
/// old hostname-based scheme). The passphrase is readable by any local process.
/// Use `Keystore::unlock()` with a user-supplied passphrase for real confidentiality.
fn machine_passphrase() -> String {
    let id_path = machine_id_path();
    let machine_id = match std::fs::read_to_string(&id_path) {
        Ok(id) => {
            let id = id.trim().to_string();
            if id.is_empty() {
                create_machine_id(&id_path)
            } else {
                id
            }
        }
        Err(_) => create_machine_id(&id_path),
    };
    format!("roboticus-machine-key:{machine_id}")
}

/// Generate a new random machine ID and persist it to disk.
fn create_machine_id(path: &std::path::Path) -> String {
    let mut bytes = [0u8; 32];
    OsRng.fill_bytes(&mut bytes);
    let id: String = bytes.iter().map(|b| format!("{b:02x}")).collect();

    if let Some(parent) = path.parent() {
        let _ = std::fs::create_dir_all(parent);
    }
    if let Err(e) = std::fs::write(path, &id) {
        tracing::error!(error = %e, path = %path.display(), "failed to write machine-id; keystore will use ephemeral ID");
    } else {
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let _ = std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600));
        }
        tracing::info!(path = %path.display(), "created new machine-id for keystore");
    }
    id
}

/// All legacy passphrase derivations to try for backward-compatible unlock.
/// Returns a list of candidate passphrases covering:
/// - Pre-rebrand (`ironclad-machine-key:`) with env vars
/// - Pre-rebrand with gethostname syscall
/// - Post-rebrand (`roboticus-machine-key:`) with env vars
/// - Post-rebrand with gethostname syscall
fn legacy_passphrases() -> Vec<String> {
    let syscall_hostname = gethostname::gethostname().to_string_lossy().into_owned();

    let env_hostname = std::env::var("HOSTNAME")
        .or_else(|_| std::env::var("HOST"))
        .unwrap_or_else(|_| "unknown-host".into());

    let username = std::env::var("USER")
        .or_else(|_| std::env::var("USERNAME"))
        .unwrap_or_else(|_| "unknown-user".into());

    let mut candidates = Vec::new();

    // Pre-rebrand prefix (the original key format)
    candidates.push(format!("ironclad-machine-key:{env_hostname}:{username}"));
    if !syscall_hostname.is_empty() && syscall_hostname != env_hostname {
        candidates.push(format!(
            "ironclad-machine-key:{syscall_hostname}:{username}"
        ));
    }

    // Post-rebrand prefix
    candidates.push(format!("roboticus-machine-key:{env_hostname}:{username}"));
    if !syscall_hostname.is_empty() && syscall_hostname != env_hostname {
        candidates.push(format!(
            "roboticus-machine-key:{syscall_hostname}:{username}"
        ));
    }

    candidates
}

// ── Shared key resolution ────────────────────────────────────────────────
//
// Canonical resolution cascade for provider API keys. Both the API
// (`roboticus-api/src/api/routes/admin.rs`) and the CLI
// (`roboticus-cli/src/cli/admin/misc/helper_model_triage.rs`) must use
// this to avoid divergence.

/// Where a provider's API key was found (or that none is available).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum KeySource {
    /// Provider runs locally and needs no key.
    NotRequired,
    /// Key found via OAuth auth_mode.
    OAuth,
    /// Key found in the keystore under the given entry name.
    Keystore(String),
    /// Key found in the given environment variable.
    EnvVar(String),
    /// No key found anywhere.
    Missing,
}

impl KeySource {
    /// Returns a `(status, source)` pair for API/dashboard display.
    pub fn status_pair(&self) -> (&'static str, &'static str) {
        match self {
            Self::NotRequired => ("not_required", "local"),
            Self::OAuth => ("configured", "oauth"),
            Self::Keystore(_) => ("configured", "keystore"),
            Self::EnvVar(_) => ("configured", "env"),
            Self::Missing => ("missing", "none"),
        }
    }

    /// Whether a usable key was found.
    pub fn is_configured(&self) -> bool {
        matches!(self, Self::Keystore(_) | Self::EnvVar(_) | Self::OAuth)
    }
}

/// Resolve the source of a provider's API key using a priority cascade:
///
/// 1. `is_local` → `NotRequired`
/// 2. `auth_mode == "oauth"` → `OAuth`
/// 3. Explicit keystore ref (`api_key_ref = "keystore:<name>"`)
/// 4. Conventional keystore name (`{provider_name}_api_key`)
/// 5. Non-empty environment variable (`api_key_env`)
/// 6. `Missing`
///
/// If the keystore is locked, keystore lookups are skipped and the cascade
/// falls through to env vars or `Missing`.
pub fn resolve_key_source(
    provider_name: &str,
    is_local: bool,
    api_key_ref: Option<&str>,
    api_key_env: Option<&str>,
    auth_mode: Option<&str>,
    keystore: &Keystore,
) -> KeySource {
    if is_local {
        return KeySource::NotRequired;
    }

    if auth_mode.is_some_and(|m| m == "oauth") {
        return KeySource::OAuth;
    }

    if let Some(ks_name) = api_key_ref.and_then(|r| r.strip_prefix("keystore:"))
        && let Some(val) = keystore.get(ks_name)
        && !val.is_empty()
    {
        return KeySource::Keystore(val);
    }

    let conventional = format!("{provider_name}_api_key");
    if let Some(val) = keystore.get(&conventional)
        && !val.is_empty()
    {
        return KeySource::Keystore(val);
    }

    if let Some(env_name) = api_key_env
        && let Ok(val) = std::env::var(env_name)
        && !val.is_empty()
    {
        return KeySource::EnvVar(val);
    }

    KeySource::Missing
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Mutex;
    use tempfile::NamedTempFile;

    /// Tests that call `unlock_machine()` or `machine_passphrase()` share the
    /// global `~/.roboticus/machine-id` file. Without serialization, parallel
    /// tests can delete/recreate the file while another test has already derived
    /// a passphrase from the old value, causing "decryption failed" errors.
    static MACHINE_ID_MUTEX: Mutex<()> = Mutex::new(());

    fn temp_path() -> PathBuf {
        let f = NamedTempFile::new().unwrap();
        let p = f.path().to_path_buf();
        drop(f);
        p
    }

    #[test]
    fn test_new_keystore_creates_empty() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        assert!(!ks.is_unlocked());

        ks.unlock("test-pass").unwrap();
        assert!(ks.is_unlocked());
        assert!(ks.list_keys().is_empty());
        assert!(path.exists());
    }

    #[test]
    fn test_set_and_get() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        ks.unlock("pass").unwrap();

        ks.set("api_key", "sk-123").unwrap();
        assert_eq!(ks.get("api_key"), Some("sk-123".into()));
        assert_eq!(ks.get("missing"), None);
    }

    #[test]
    fn test_persistence() {
        let path = temp_path();

        {
            let ks = Keystore::new(&path);
            ks.unlock("my-pass").unwrap();
            ks.set("secret", "value42").unwrap();
        }

        {
            let ks = Keystore::new(&path);
            assert!(!ks.is_unlocked());
            ks.unlock("my-pass").unwrap();
            assert_eq!(ks.get("secret"), Some("value42".into()));
        }
    }

    #[test]
    fn test_wrong_passphrase() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        ks.unlock("correct").unwrap();
        ks.set("key", "val").unwrap();
        drop(ks);

        let ks2 = Keystore::new(&path);
        let result = ks2.unlock("wrong");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("decryption"));
    }

    #[test]
    fn test_list_keys() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        ks.unlock("pass").unwrap();

        ks.set("alpha", "1").unwrap();
        ks.set("beta", "2").unwrap();
        ks.set("gamma", "3").unwrap();

        let mut keys = ks.list_keys();
        keys.sort();
        assert_eq!(keys, vec!["alpha", "beta", "gamma"]);
    }

    #[test]
    fn test_remove() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        ks.unlock("pass").unwrap();

        ks.set("keep", "a").unwrap();
        ks.set("discard", "b").unwrap();

        assert!(ks.remove("discard").unwrap());
        assert!(!ks.remove("discard").unwrap());
        assert_eq!(ks.get("discard"), None);
        assert_eq!(ks.get("keep"), Some("a".into()));

        drop(ks);
        let ks2 = Keystore::new(&path);
        ks2.unlock("pass").unwrap();
        assert_eq!(ks2.get("discard"), None);
        assert_eq!(ks2.get("keep"), Some("a".into()));
    }

    #[test]
    fn test_import() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        ks.unlock("pass").unwrap();

        let mut batch = HashMap::new();
        batch.insert("k1".into(), "v1".into());
        batch.insert("k2".into(), "v2".into());
        batch.insert("k3".into(), "v3".into());

        let count = ks.import(batch).unwrap();
        assert_eq!(count, 3);
        assert_eq!(ks.get("k1"), Some("v1".into()));
        assert_eq!(ks.get("k2"), Some("v2".into()));
        assert_eq!(ks.get("k3"), Some("v3".into()));
    }

    #[test]
    fn test_machine_key() {
        let _lock = MACHINE_ID_MUTEX.lock().unwrap();
        let path = temp_path();
        let ks = Keystore::new(&path);
        ks.unlock_machine().unwrap();
        ks.set("service_key", "abc").unwrap();
        drop(ks);

        let ks2 = Keystore::new(&path);
        ks2.unlock_machine().unwrap();
        assert_eq!(ks2.get("service_key"), Some("abc".into()));
    }

    #[test]
    fn test_get_refreshes_entries_after_external_write() {
        let _lock = MACHINE_ID_MUTEX.lock().unwrap();
        let path = temp_path();
        let ks_a = Keystore::new(&path);
        ks_a.unlock_machine().unwrap();
        ks_a.set("openai_api_key", "old-value").unwrap();
        assert_eq!(ks_a.get("openai_api_key"), Some("old-value".into()));

        // Simulate a second process mutating the same keystore file.
        let ks_b = Keystore::new(&path);
        ks_b.unlock_machine().unwrap();
        ks_b.set("openai_api_key", "new-value").unwrap();

        // Without a refresh-on-read policy this can stay stale in ks_a.
        assert_eq!(ks_a.get("openai_api_key"), Some("new-value".into()));
    }

    #[test]
    fn test_lock_clears_memory() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        ks.unlock("pass").unwrap();
        ks.set("secret", "hidden").unwrap();
        assert!(ks.is_unlocked());

        ks.lock();

        assert!(!ks.is_unlocked());
        assert_eq!(ks.get("secret"), None);
        assert!(ks.list_keys().is_empty());
    }

    #[test]
    fn test_rekey() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        ks.unlock("old-pass").unwrap();
        ks.set("data", "preserved").unwrap();

        ks.rekey("new-pass").unwrap();
        drop(ks);

        let ks2 = Keystore::new(&path);
        assert!(ks2.unlock("old-pass").is_err());
        ks2.unlock("new-pass").unwrap();
        assert_eq!(ks2.get("data"), Some("preserved".into()));
    }

    #[test]
    fn test_keystore_mutations_are_audited() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        ks.unlock("pass").unwrap();
        ks.set("telegram_bot_token", "secret").unwrap();
        assert!(ks.remove("telegram_bot_token").unwrap());
        ks.rekey("new-pass").unwrap();

        let audit_path = path.with_extension("audit.log");
        let audit = std::fs::read_to_string(audit_path).unwrap();
        assert!(audit.contains("\"operation\":\"initialize\""));
        assert!(audit.contains("\"operation\":\"set\""));
        assert!(audit.contains("\"operation\":\"remove\""));
        assert!(audit.contains("\"operation\":\"rekey\""));
        // Key names are redacted: only first 3 chars visible, followed by ***
        assert!(audit.contains("\"key\":\"tel***\""));
        assert!(!audit.contains("telegram_bot_token"));
        assert!(!audit.contains("secret"));
    }

    #[test]
    fn test_default_path() {
        let path = Keystore::default_path();
        assert!(path.to_str().unwrap().contains("keystore.enc"));
        assert!(path.to_str().unwrap().contains(".roboticus"));
    }

    #[test]
    fn test_set_on_locked_keystore_fails() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        // Don't unlock
        let result = ks.set("key", "value");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("locked"));
    }

    #[test]
    fn test_remove_on_locked_keystore_fails() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        let result = ks.remove("key");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("locked"));
    }

    #[test]
    fn test_import_on_locked_keystore_fails() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        let result = ks.import(HashMap::new());
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("locked"));
    }

    #[test]
    fn test_rekey_on_locked_keystore_fails() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        let result = ks.rekey("new-pass");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("locked"));
    }

    #[test]
    fn test_get_on_locked_keystore_returns_none() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        assert_eq!(ks.get("anything"), None);
    }

    #[test]
    fn test_list_keys_on_locked_keystore_returns_empty() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        assert!(ks.list_keys().is_empty());
    }

    #[test]
    fn test_corrupt_keystore_file() {
        let path = temp_path();
        // Write too-short data (less than SALT_LEN + NONCE_LEN + 1)
        std::fs::write(&path, b"short").unwrap();
        let ks = Keystore::new(&path);
        let result = ks.unlock("pass");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("corrupt"));
    }

    #[test]
    fn test_set_overwrites_existing_key() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        ks.unlock("pass").unwrap();

        ks.set("key", "first").unwrap();
        assert_eq!(ks.get("key"), Some("first".into()));

        ks.set("key", "second").unwrap();
        assert_eq!(ks.get("key"), Some("second".into()));
    }

    #[cfg(unix)]
    #[test]
    fn test_set_rolls_back_on_save_failure() {
        use std::os::unix::fs::PermissionsExt;
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("keystore.enc");
        let ks = Keystore::new(&path);
        ks.unlock("pass").unwrap();
        ks.set("stable", "1").unwrap();

        let mut perms = std::fs::metadata(dir.path()).unwrap().permissions();
        perms.set_mode(0o500);
        std::fs::set_permissions(dir.path(), perms).unwrap();

        let res = ks.set("transient", "2");
        assert!(res.is_err());
        assert_eq!(ks.get("stable"), Some("1".into()));
        assert_eq!(ks.get("transient"), None);
        let audit = std::fs::read_to_string(path.with_extension("audit.log")).unwrap();
        assert!(audit.contains("\"operation\":\"set\""));
        assert!(audit.contains("\"rolled_back\":true"));

        let mut restore = std::fs::metadata(dir.path()).unwrap().permissions();
        restore.set_mode(0o700);
        std::fs::set_permissions(dir.path(), restore).unwrap();
    }

    #[test]
    fn test_import_audit_entry() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        ks.unlock("pass").unwrap();

        let mut batch = HashMap::new();
        batch.insert("imported_key".into(), "imported_value".into());
        ks.import(batch).unwrap();

        let audit_path = path.with_extension("audit.log");
        let audit = std::fs::read_to_string(audit_path).unwrap();
        assert!(audit.contains("\"operation\":\"import\""));
    }

    #[test]
    fn redact_key_name_short_keys() {
        assert_eq!(redact_key_name("ab"), "ab***");
        assert_eq!(redact_key_name("a"), "a***");
        assert_eq!(redact_key_name(""), "***");
    }

    #[test]
    fn redact_key_name_long_keys() {
        assert_eq!(redact_key_name("telegram_bot_token"), "tel***");
        assert_eq!(redact_key_name("abc"), "abc***");
    }

    /// Direct tests to a temp directory so they never touch the production
    /// machine-id at `~/.roboticus/machine-id`.
    ///
    /// SAFETY: set_var/remove_var are unsafe in Rust 2024 because they're
    /// not thread-safe. We hold MACHINE_ID_MUTEX so no concurrent access.
    fn with_test_machine_id_dir<F: FnOnce()>(f: F) {
        let dir = tempfile::tempdir().unwrap();
        unsafe { std::env::set_var("ROBOTICUS_TEST_MACHINE_ID_DIR", dir.path()) };
        f();
        unsafe { std::env::remove_var("ROBOTICUS_TEST_MACHINE_ID_DIR") };
    }

    #[test]
    fn machine_passphrase_is_deterministic() {
        let _lock = MACHINE_ID_MUTEX.lock().unwrap();
        with_test_machine_id_dir(|| {
            let p1 = machine_passphrase();
            let p2 = machine_passphrase();
            assert_eq!(p1, p2);
            assert!(p1.starts_with("roboticus-machine-key:"));
        });
    }

    #[test]
    fn machine_id_persists_across_calls() {
        let _lock = MACHINE_ID_MUTEX.lock().unwrap();
        with_test_machine_id_dir(|| {
            // machine_passphrase creates a machine-id on first call;
            // subsequent calls must return the same value.
            let p1 = machine_passphrase();
            let p2 = machine_passphrase();
            assert_eq!(p1, p2);
            let id_path = machine_id_path();
            assert!(id_path.exists());
            let contents = std::fs::read_to_string(&id_path).unwrap();
            assert_eq!(contents.trim().len(), 64); // 32 bytes = 64 hex chars
        });
    }

    #[test]
    fn legacy_passphrases_include_both_prefixes() {
        let candidates = legacy_passphrases();
        assert!(!candidates.is_empty());
        // Must include at least the ironclad prefix (pre-rebrand)
        assert!(
            candidates
                .iter()
                .any(|p| p.starts_with("ironclad-machine-key:"))
        );
        // Must include at least the roboticus prefix (post-rebrand)
        assert!(
            candidates
                .iter()
                .any(|p| p.starts_with("roboticus-machine-key:"))
        );
    }

    #[test]
    fn unlock_machine_migrates_legacy_keystore() {
        let _lock = MACHINE_ID_MUTEX.lock().unwrap();
        with_test_machine_id_dir(|| {
            // Simulate a keystore created with a legacy passphrase.
            let path = temp_path();
            let candidates = legacy_passphrases();
            let legacy_pass = &candidates[0]; // first legacy candidate
            let ks = Keystore::new(&path);
            ks.unlock(legacy_pass).unwrap();
            ks.set("secret", "migrated").unwrap();
            drop(ks);

            // unlock_machine should find it via legacy fallback and auto-rekey.
            let ks2 = Keystore::new(&path);
            ks2.unlock_machine().unwrap();
            assert_eq!(ks2.get("secret"), Some("migrated".into()));

            // After migration, the primary (machine-id) passphrase should work directly.
            drop(ks2);
            let primary = machine_passphrase();
            let ks3 = Keystore::new(&path);
            ks3.unlock(&primary).unwrap();
            assert_eq!(ks3.get("secret"), Some("migrated".into()));
        });
    }

    #[test]
    fn unlock_machine_recovers_pre_rebrand_keystore() {
        let _lock = MACHINE_ID_MUTEX.lock().unwrap();
        // Simulate a keystore created with the ironclad prefix (pre-rebrand).
        let path = temp_path();
        let hostname = std::env::var("HOST")
            .or_else(|_| std::env::var("HOSTNAME"))
            .unwrap_or_else(|_| "unknown-host".into());
        let username = std::env::var("USER")
            .or_else(|_| std::env::var("USERNAME"))
            .unwrap_or_else(|_| "unknown-user".into());
        let old_pass = format!("ironclad-machine-key:{hostname}:{username}");

        let ks = Keystore::new(&path);
        ks.unlock(&old_pass).unwrap();
        ks.set("discord_token", "abc123").unwrap();
        drop(ks);

        // unlock_machine must recover it.
        let ks2 = Keystore::new(&path);
        ks2.unlock_machine().unwrap();
        assert_eq!(ks2.get("discord_token"), Some("abc123".into()));
    }

    #[test]
    fn lock_or_recover_works_on_clean_mutex() {
        let m = Mutex::new(42);
        let guard = lock_or_recover(&m);
        assert_eq!(*guard, 42);
    }

    #[test]
    fn audit_log_path_derives_from_keystore_path() {
        let ks = Keystore::new("/tmp/test.enc");
        assert_eq!(ks.audit_log_path(), PathBuf::from("/tmp/test.audit.log"));
    }

    #[test]
    fn concurrent_set_and_rekey_no_deadlock() {
        let path = temp_path();
        let ks = Keystore::new(&path);
        ks.unlock("pass").unwrap();
        const ITERATIONS: usize = 10;

        std::thread::scope(|s| {
            let ks1 = ks.clone();
            let ks2 = ks.clone();

            let h1 = s.spawn(move || {
                for i in 0..ITERATIONS {
                    ks1.set(&format!("key-{i}"), &format!("val-{i}")).unwrap();
                }
            });
            let h2 = s.spawn(move || {
                for _ in 0..ITERATIONS {
                    ks2.rekey("pass").unwrap();
                }
            });

            // Both threads must complete (no deadlock) while keeping this test fast
            // enough for CI and local release gates.
            h1.join().unwrap();
            h2.join().unwrap();
        });
    }
}