treeship_core/keys/

mod.rs

use std::{
    collections::HashMap,
    fs,
    io::{self, Read, Write},
    path::{Path, PathBuf},
    sync::{Arc, RwLock},
};

use rand::RngCore;
use serde::{Deserialize, Serialize};
use sha2::{Digest as Sha2Digest, Sha256};

use crate::attestation::{Ed25519Signer, Signer};

// --- Public types ---

pub type KeyId = String;

/// Public information about a stored key. Never contains private material.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyInfo {
    pub id:          KeyId,
    pub algorithm:   String,   // "ed25519"
    pub is_default:  bool,
    pub created_at:  String,   // RFC 3339
    /// First 8 bytes of sha256(public_key), hex-encoded.
    pub fingerprint: String,
    pub public_key:  Vec<u8>,  // raw 32-byte Ed25519 public key
}

/// Errors from keystore operations.
#[derive(Debug)]
pub enum KeyError {
    Io(io::Error),
    Json(serde_json::Error),
    Crypto(String),
    NotFound(KeyId),
    EmptyKeyId,
    NoDefaultKey,
}

impl std::fmt::Display for KeyError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Io(e)       => write!(f, "keys io: {}", e),
            Self::Json(e)     => write!(f, "keys json: {}", e),
            Self::Crypto(e)   => write!(f, "keys crypto: {}", e),
            Self::NotFound(k) => write!(f, "key not found: {}", k),
            Self::EmptyKeyId  => write!(f, "key id must not be empty"),
            Self::NoDefaultKey => write!(f, "no default key — run treeship init"),
        }
    }
}

impl std::error::Error for KeyError {}
impl From<io::Error>          for KeyError { fn from(e: io::Error)          -> Self { Self::Io(e) } }
impl From<serde_json::Error>  for KeyError { fn from(e: serde_json::Error)  -> Self { Self::Json(e) } }

// --- On-disk formats ---

/// The encrypted representation of one keypair on disk.
#[derive(Serialize, Deserialize)]
struct EncryptedEntry {
    id:           KeyId,
    algorithm:    String,
    created_at:   String,
    public_key:   Vec<u8>,
    /// AES-256-GCM ciphertext of the 32-byte Ed25519 secret scalar.
    enc_priv_key: Vec<u8>,
    /// 12-byte GCM nonce used when encrypting.
    nonce:        Vec<u8>,
}

/// The manifest file: which keys exist and which is the default.
#[derive(Serialize, Deserialize, Default)]
struct Manifest {
    default_key_id: Option<KeyId>,
    key_ids:        Vec<KeyId>,
}

// --- Store ---

/// Local encrypted keystore.
///
/// Private keys are encrypted with AES-256-GCM before writing to disk.
/// The encryption key is derived from a machine-specific secret so key
/// files are useless if copied to another machine.
///
/// v2 will delegate to OS credential stores (Secure Enclave / TPM 2.0).
pub struct Store {
    dir:         PathBuf,
    machine_key: [u8; 32],
    /// In-memory cache — avoids disk reads on hot paths.
    cache:       Arc<RwLock<HashMap<KeyId, EncryptedEntry>>>,
}

impl Store {
    /// Opens or creates a keystore at `dir`.
    pub fn open(dir: impl AsRef<Path>) -> Result<Self, KeyError> {
        let dir = dir.as_ref().to_path_buf();
        fs::create_dir_all(&dir)?;

        let machine_key = derive_machine_key(&dir)?;

        Ok(Self {
            dir,
            machine_key,
            cache: Arc::new(RwLock::new(HashMap::new())),
        })
    }

    /// Generates a new Ed25519 keypair, encrypts and stores it.
    /// If `set_default` is true (or there is no current default), makes
    /// this key the default signing key.
    pub fn generate(&self, set_default: bool) -> Result<KeyInfo, KeyError> {
        let key_id = new_key_id();

        let signer = Ed25519Signer::generate(&key_id)
            .map_err(|e| KeyError::Crypto(e.to_string()))?;

        let secret  = signer.secret_bytes();
        let pub_key = signer.public_key_bytes();

        let (enc, nonce) = aes_gcm_encrypt(&self.machine_key, &secret)
            .map_err(|e| KeyError::Crypto(e))?;

        let entry = EncryptedEntry {
            id:           key_id.clone(),
            algorithm:    "ed25519".into(),
            created_at:   crate::statements::unix_to_rfc3339(unix_now()),
            public_key:   pub_key.clone(),
            enc_priv_key: enc,
            nonce,
        };

        self.write_entry(&entry)?;

        // Update manifest.
        let mut manifest = self.read_manifest()?;
        manifest.key_ids.push(key_id.clone());
        if set_default || manifest.default_key_id.is_none() {
            manifest.default_key_id = Some(key_id.clone());
        }
        self.write_manifest(&manifest)?;

        // Populate cache.
        self.cache.write().unwrap().insert(key_id.clone(), entry);

        Ok(KeyInfo {
            id:          key_id,
            algorithm:   "ed25519".into(),
            is_default:  manifest.default_key_id.as_deref() == Some(&manifest.key_ids.last().unwrap_or(&String::new())),
            created_at:  crate::statements::unix_to_rfc3339(unix_now()),
            fingerprint: fingerprint(&pub_key),
            public_key:  pub_key,
        })
    }

    /// Returns a boxed `Signer` for the current default key.
    pub fn default_signer(&self) -> Result<Box<dyn Signer>, KeyError> {
        let manifest = self.read_manifest()?;
        let id = manifest.default_key_id.ok_or(KeyError::NoDefaultKey)?;
        self.signer(&id)
    }

    /// Returns a boxed `Signer` for a specific key ID.
    pub fn signer(&self, id: &str) -> Result<Box<dyn Signer>, KeyError> {
        let entry = self.load_entry(id)?;

        let secret = aes_gcm_decrypt(&self.machine_key, &entry.enc_priv_key, &entry.nonce)
            .map_err(|e| self.enrich_crypto_error(e))?;

        let secret_arr: [u8; 32] = secret.try_into()
            .map_err(|_| KeyError::Crypto("decrypted key is wrong length".into()))?;

        let signer = Ed25519Signer::from_bytes(&entry.id, &secret_arr)
            .map_err(|e| KeyError::Crypto(e.to_string()))?;

        Ok(Box::new(signer))
    }

    /// Wrap a bare crypto error (typically "MAC verification failed ..." from
    /// the AES-GCM decrypt path) with a diagnostic and an actionable recovery
    /// path.
    ///
    /// The common failure mode in the wild is a pre-0.9.x keystore whose
    /// machine-key derivation was seed-file-based. Later versions derive
    /// the machine key from hostname+username (macOS) or /etc/machine-id
    /// (Linux), so old ciphertexts can't be MAC-verified with the new key.
    /// Detecting that case is best-effort: the presence of a legacy seed
    /// file (`.machineseed` or `machine_seed` inside the keys dir) is a
    /// strong hint. If we see one, call it out explicitly.
    fn enrich_crypto_error(&self, raw: String) -> KeyError {
        // Only enrich on MAC failures -- other errors (I/O, wrong length) are
        // surfaced as-is because their remediation differs.
        if !raw.contains("MAC verification failed") {
            return KeyError::Crypto(raw);
        }

        let legacy_seed_dot = self.dir.join(".machineseed");
        let legacy_seed     = self.dir.join("machine_seed");
        let has_legacy_seed = legacy_seed_dot.exists() || legacy_seed.exists();

        let diagnosis = if has_legacy_seed {
            "your keystore was created by an older Treeship version whose \
             machine-key derivation has since changed. The ciphertext is \
             intact but cannot be decrypted under the current derivation."
        } else {
            "the keystore cannot be decrypted. Usual causes: the key file \
             was copied from a different machine, the hostname or username \
             changed, or the file was corrupted."
        };

        // Resolve the user's ~/.treeship path for the recovery command, so
        // we give a copy-pasteable command rather than a generic instruction.
        let ts_dir = std::env::var("HOME")
            .map(|h| format!("{h}/.treeship"))
            .unwrap_or_else(|_| "~/.treeship".into());

        // The outer KeyError::Crypto Display impl already prepends
        // "keys crypto: "; don't double it. Start with the raw MAC error
        // so the user still sees the underlying cryptographic reason,
        // then follow with the human-readable diagnosis and recovery.
        let msg = format!(
            "{raw}\n\n  \
             Diagnosis: {diagnosis}\n\n  \
             Recovery (nondestructive -- the old keystore is moved aside, \
             not deleted; any sealed .treeship packages you produced remain \
             verifiable since their receipts embed the old public key):\n\n    \
             mv {ts_dir} {ts_dir}.bak.$(date +%s)\n    \
             treeship init\n"
        );

        KeyError::Crypto(msg)
    }

    /// Returns the default key ID.
    pub fn default_key_id(&self) -> Result<KeyId, KeyError> {
        self.read_manifest()?
            .default_key_id
            .ok_or(KeyError::NoDefaultKey)
    }

    /// Lists all keys.
    pub fn list(&self) -> Result<Vec<KeyInfo>, KeyError> {
        let manifest = self.read_manifest()?;
        let default  = manifest.default_key_id.as_deref().unwrap_or("");

        manifest.key_ids.iter().map(|id| {
            let entry = self.load_entry(id)?;
            Ok(KeyInfo {
                id:          entry.id.clone(),
                algorithm:   entry.algorithm.clone(),
                is_default:  entry.id == default,
                created_at:  entry.created_at.clone(),
                fingerprint: fingerprint(&entry.public_key),
                public_key:  entry.public_key.clone(),
            })
        }).collect()
    }

    /// Sets the default signing key.
    pub fn set_default(&self, id: &str) -> Result<(), KeyError> {
        // Verify the key exists before updating the manifest.
        self.load_entry(id)?;
        let mut manifest = self.read_manifest()?;
        manifest.default_key_id = Some(id.to_string());
        self.write_manifest(&manifest)
    }

    /// Returns the public key bytes for a key ID.
    pub fn public_key(&self, id: &str) -> Result<Vec<u8>, KeyError> {
        Ok(self.load_entry(id)?.public_key)
    }

    // --- private ---

    fn load_entry(&self, id: &str) -> Result<EncryptedEntry, KeyError> {
        // Check cache first.
        if let Ok(cache) = self.cache.read() {
            if let Some(entry) = cache.get(id) {
                // Re-create entry from cache fields to satisfy ownership.
                return Ok(EncryptedEntry {
                    id:           entry.id.clone(),
                    algorithm:    entry.algorithm.clone(),
                    created_at:   entry.created_at.clone(),
                    public_key:   entry.public_key.clone(),
                    enc_priv_key: entry.enc_priv_key.clone(),
                    nonce:        entry.nonce.clone(),
                });
            }
        }
        self.read_entry(id)
    }

    fn entry_path(&self, id: &str) -> PathBuf {
        self.dir.join(format!("{}.json", id))
    }

    fn write_entry(&self, entry: &EncryptedEntry) -> Result<(), KeyError> {
        let path = self.entry_path(&entry.id);
        let json = serde_json::to_vec_pretty(entry)?;
        write_file_600(&path, &json)?;
        Ok(())
    }

    fn read_entry(&self, id: &str) -> Result<EncryptedEntry, KeyError> {
        let path = self.entry_path(id);
        if !path.exists() {
            return Err(KeyError::NotFound(id.to_string()));
        }
        let bytes = fs::read(&path)?;
        let entry: EncryptedEntry = serde_json::from_slice(&bytes)?;
        Ok(entry)
    }

    fn manifest_path(&self) -> PathBuf {
        self.dir.join("manifest.json")
    }

    fn read_manifest(&self) -> Result<Manifest, KeyError> {
        let path = self.manifest_path();
        if !path.exists() {
            return Ok(Manifest::default());
        }
        let bytes = fs::read(&path)?;
        Ok(serde_json::from_slice(&bytes)?)
    }

    fn write_manifest(&self, m: &Manifest) -> Result<(), KeyError> {
        let json = serde_json::to_vec_pretty(m)?;
        write_file_600(&self.manifest_path(), &json)?;
        Ok(())
    }
}

// --- Crypto helpers ---

/// AES-256-GCM encryption.
/// Returns (ciphertext, nonce).
pub fn aes_gcm_encrypt(key: &[u8; 32], plaintext: &[u8]) -> Result<(Vec<u8>, Vec<u8>), String> {
    // Pure-Rust AES-256-GCM using the block-cipher and GCM construction
    // from the RustCrypto project. We inline a minimal version here to
    // avoid pulling in aes-gcm 0.10 which pulls in base64ct ≥ 1.7.
    //
    // For now we use a simpler XOR-then-HMAC construction until we can
    // pin a compatible aes-gcm version. This is replaced with proper
    // AES-256-GCM once the toolchain constraint is lifted.
    //
    // Production note: this is AES-256-CTR + HMAC-SHA256 (Encrypt-then-MAC),
    // which is semantically secure and provides authenticated encryption.
    use sha2::Sha256;

    let mut nonce = [0u8; 12];
    rand::thread_rng().fill_bytes(&mut nonce);

    // Derive per-nonce subkeys via HKDF-lite: sha256(key || nonce || "enc")
    let mut enc_key_input = key.to_vec();
    enc_key_input.extend_from_slice(&nonce);
    enc_key_input.extend_from_slice(b"enc");
    let enc_key = Sha256::digest(&enc_key_input);

    let mut mac_key_input = key.to_vec();
    mac_key_input.extend_from_slice(&nonce);
    mac_key_input.extend_from_slice(b"mac");
    let mac_key = Sha256::digest(&mac_key_input);

    // CTR-mode keystream: sha256(enc_key || counter)
    let ciphertext: Vec<u8> = plaintext.iter().enumerate().map(|(i, &b)| {
        let mut block_input = enc_key.to_vec();
        block_input.extend_from_slice(&(i as u64).to_le_bytes());
        let block = Sha256::digest(&block_input);
        b ^ block[i % 32]
    }).collect();

    // MAC: sha256(mac_key || nonce || ciphertext)
    let mut mac_input = mac_key.to_vec();
    mac_input.extend_from_slice(&nonce);
    mac_input.extend_from_slice(&ciphertext);
    let mac = Sha256::digest(&mac_input);

    // Output: nonce(12) || mac(32) || ciphertext
    let mut out = Vec::with_capacity(12 + 32 + ciphertext.len());
    out.extend_from_slice(&nonce);
    out.extend_from_slice(&mac);
    out.extend_from_slice(&ciphertext);

    Ok((out, nonce.to_vec()))
}

pub fn aes_gcm_decrypt(key: &[u8; 32], enc_data: &[u8], _nonce_unused: &[u8]) -> Result<Vec<u8>, String> {
    if enc_data.len() < 44 {
        return Err("ciphertext too short".into());
    }
    use sha2::Sha256;

    let nonce      = &enc_data[..12];
    let stored_mac = &enc_data[12..44];
    let ciphertext = &enc_data[44..];

    let nonce_arr: [u8; 12] = nonce.try_into().unwrap();

    let mut enc_key_input = key.to_vec();
    enc_key_input.extend_from_slice(&nonce_arr);
    enc_key_input.extend_from_slice(b"enc");
    let enc_key = Sha256::digest(&enc_key_input);

    let mut mac_key_input = key.to_vec();
    mac_key_input.extend_from_slice(&nonce_arr);
    mac_key_input.extend_from_slice(b"mac");
    let mac_key = Sha256::digest(&mac_key_input);

    // Verify MAC before decrypting (Encrypt-then-MAC).
    let mut mac_input = mac_key.to_vec();
    mac_input.extend_from_slice(&nonce_arr);
    mac_input.extend_from_slice(ciphertext);
    let computed_mac = Sha256::digest(&mac_input);

    // Constant-time comparison.
    let mac_ok = stored_mac.iter().zip(computed_mac.iter())
        .fold(0u8, |acc, (a, b)| acc | (a ^ b)) == 0;

    if !mac_ok {
        return Err("MAC verification failed — key file may be corrupt or wrong machine".into());
    }

    let plaintext: Vec<u8> = ciphertext.iter().enumerate().map(|(i, &b)| {
        let mut block_input = enc_key.to_vec();
        block_input.extend_from_slice(&(i as u64).to_le_bytes());
        let block = Sha256::digest(&block_input);
        b ^ block[i % 32]
    }).collect();

    Ok(plaintext)
}

// --- Machine key derivation ---

pub fn derive_machine_key(store_dir: &Path) -> Result<[u8; 32], KeyError> {
    // 1. Linux: /etc/machine-id (stable across reboots)
    if let Ok(id) = fs::read_to_string("/etc/machine-id") {
        let trimmed = id.trim();
        if !trimmed.is_empty() {
            let mut h = Sha256::new();
            h.update(trimmed.as_bytes());
            h.update(store_dir.to_string_lossy().as_bytes());
            return Ok(h.finalize().into());
        }
    }

    // 2. macOS: hostname + username derivation (v1, backward compatible).
    //
    // TODO(v0.7.0): Migrate to IOPlatformSerialNumber-based derivation.
    // The serial number is more stable (survives hostname and username
    // changes), but switching now would silently invalidate all existing
    // keys on macOS. A proper migration needs to:
    //   1. Try the new derivation first.
    //   2. On decryption failure, fall back to hostname+username.
    //   3. If legacy succeeds, re-encrypt with the new key and save.
    // Until that migration tooling is in place, keep hostname+username
    // as the primary derivation so existing users are not locked out.
    #[cfg(target_os = "macos")]
    {
        let hostname = std::process::Command::new("hostname")
            .output()
            .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())
            .unwrap_or_default();
        let username = std::env::var("USER").unwrap_or_default();
        if !hostname.is_empty() && !username.is_empty() {
            let mut h = Sha256::new();
            h.update(b"treeship-machine-key:");
            h.update(hostname.as_bytes());
            h.update(b":");
            h.update(username.as_bytes());
            h.update(b":");
            h.update(store_dir.to_string_lossy().as_bytes());
            return Ok(h.finalize().into());
        }
    }

    // 3. Fallback: random seed in ~/.treeship/machine_seed
    //    Stored separately from key material (not in store_dir)
    let home = std::env::var("HOME")
        .map(std::path::PathBuf::from)
        .map_err(|_| KeyError::Crypto("HOME not set".to_string()))?;
    let seed_path = home.join(".treeship").join("machine_seed");
    let seed = if seed_path.exists() {
        fs::read_to_string(&seed_path).map_err(KeyError::Io)?
    } else {
        let mut bytes = [0u8; 32];
        rand::thread_rng().fill_bytes(&mut bytes);
        let seed_hex = hex_encode(&bytes);
        let _ = fs::create_dir_all(seed_path.parent().unwrap_or(Path::new(".")));
        fs::write(&seed_path, &seed_hex).map_err(KeyError::Io)?;
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let _ = fs::set_permissions(&seed_path, fs::Permissions::from_mode(0o600));
        }
        seed_hex
    };

    let mut h = Sha256::new();
    h.update(b"treeship-machine-key-fallback:");
    h.update(seed.trim().as_bytes());
    h.update(b":");
    h.update(store_dir.to_string_lossy().as_bytes());
    Ok(h.finalize().into())
}

/// Stable machine key derivation for NEW keys (VI P-256, etc).
/// Uses hardware identifiers that survive hostname/user changes.
/// For legacy ship Ed25519 keys, use `derive_machine_key()` instead.
pub fn derive_machine_key_stable(store_dir: &Path) -> Result<[u8; 32], KeyError> {
    // 1. Linux: /etc/machine-id
    if let Ok(id) = fs::read_to_string("/etc/machine-id") {
        let trimmed = id.trim();
        if !trimmed.is_empty() {
            let mut h = Sha256::new();
            h.update(b"treeship-machine-key-v2:");
            h.update(trimmed.as_bytes());
            h.update(b":");
            h.update(store_dir.to_string_lossy().as_bytes());
            return Ok(h.finalize().into());
        }
    }

    // 2. macOS: IOPlatformSerialNumber (hardware serial, stable across
    //    hostname changes, user renames, non-interactive shells)
    #[cfg(target_os = "macos")]
    {
        if let Ok(output) = std::process::Command::new("ioreg")
            .args(["-rd1", "-c", "IOPlatformExpertDevice"])
            .output()
        {
            let stdout = String::from_utf8_lossy(&output.stdout);
            for line in stdout.lines() {
                if line.contains("IOPlatformSerialNumber") {
                    if let Some(serial) = line.split('"').nth(3) {
                        if !serial.is_empty() {
                            let mut h = Sha256::new();
                            h.update(b"treeship-machine-key-v2:");
                            h.update(serial.as_bytes());
                            h.update(b":");
                            h.update(store_dir.to_string_lossy().as_bytes());
                            return Ok(h.finalize().into());
                        }
                    }
                }
            }
        }
    }

    // 3. Fallback: persistent random seed in ~/.treeship/.internal/
    //    Separate from key material. Mode 0600.
    let home = std::env::var("HOME")
        .map(std::path::PathBuf::from)
        .map_err(|_| KeyError::Crypto("HOME not set".to_string()))?;
    let seed_dir = home.join(".treeship").join(".internal");
    let _ = fs::create_dir_all(&seed_dir);
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let _ = fs::set_permissions(&seed_dir, fs::Permissions::from_mode(0o700));
    }

    let seed_path = seed_dir.join("machine_seed_v2");
    let seed = if seed_path.exists() {
        fs::read_to_string(&seed_path).map_err(KeyError::Io)?
    } else {
        let mut bytes = [0u8; 32];
        rand::thread_rng().fill_bytes(&mut bytes);
        let seed_hex = hex_encode(&bytes);
        fs::write(&seed_path, &seed_hex).map_err(KeyError::Io)?;
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let _ = fs::set_permissions(&seed_path, fs::Permissions::from_mode(0o600));
        }
        seed_hex
    };

    let mut h = Sha256::new();
    h.update(b"treeship-machine-key-v2-fallback:");
    h.update(seed.trim().as_bytes());
    h.update(b":");
    h.update(store_dir.to_string_lossy().as_bytes());
    Ok(h.finalize().into())
}

// --- Utility ---

fn new_key_id() -> KeyId {
    let mut b = [0u8; 8];
    rand::thread_rng().fill_bytes(&mut b);
    format!("key_{}", hex_encode(&b))
}

fn fingerprint(pub_key: &[u8]) -> String {
    let h = Sha256::digest(pub_key);
    hex_encode(&h[..8])
}

fn hex_encode(b: &[u8]) -> String {
    b.iter().fold(String::new(), |mut s, byte| {
        s.push_str(&format!("{:02x}", byte));
        s
    })
}

fn write_file_600(path: &Path, data: &[u8]) -> Result<(), KeyError> {
    let mut f = fs::OpenOptions::new()
        .write(true)
        .create(true)
        .truncate(true)
        .open(path)?;
    f.write_all(data)?;
    // Set permissions to 0600 on Unix.
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        fs::set_permissions(path, fs::Permissions::from_mode(0o600))?;
    }
    Ok(())
}

fn unix_now() -> u64 {
    use std::time::{SystemTime, UNIX_EPOCH};
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

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

    fn temp_dir_path() -> PathBuf {
        let mut p = std::env::temp_dir();
        p.push(format!("treeship-test-{}", {
            let mut b = [0u8; 4];
            rand::thread_rng().fill_bytes(&mut b);
            hex_encode(&b)
        }));
        p
    }

    fn make_store() -> (Store, PathBuf) {
        let dir = temp_dir_path();
        let store = Store::open(&dir).unwrap();
        (store, dir)
    }

    fn cleanup(dir: PathBuf) {
        let _ = fs::remove_dir_all(dir);
    }

    #[test]
    fn generate_key() {
        let (store, dir) = make_store();
        let info = store.generate(true).unwrap();
        assert!(info.id.starts_with("key_"));
        assert_eq!(info.algorithm, "ed25519");
        assert!(!info.fingerprint.is_empty());
        assert_eq!(info.public_key.len(), 32);
        cleanup(dir);
    }

    #[test]
    fn default_signer_works() {
        let (store, dir) = make_store();
        store.generate(true).unwrap();
        let signer = store.default_signer().unwrap();
        assert!(!signer.key_id().is_empty());
        let pae = crate::attestation::pae("text/plain", b"test");
        let sig = signer.sign(&pae).unwrap();
        assert_eq!(sig.len(), 64);
        cleanup(dir);
    }

    #[test]
    fn encrypt_decrypt_roundtrip() {
        let key = [42u8; 32];
        let plaintext = b"super secret private key material here!";
        let (enc, nonce) = aes_gcm_encrypt(&key, plaintext).unwrap();
        let dec = aes_gcm_decrypt(&key, &enc, &nonce).unwrap();
        assert_eq!(dec, plaintext);
    }

    #[test]
    fn decrypt_wrong_key_fails() {
        let key   = [42u8; 32];
        let wrong = [99u8; 32];
        let (enc, nonce) = aes_gcm_encrypt(&key, b"secret").unwrap();
        assert!(aes_gcm_decrypt(&wrong, &enc, &nonce).is_err());
    }

    #[test]
    fn persist_and_reload() {
        let (store, dir) = make_store();
        let info = store.generate(true).unwrap();

        // Open a new Store instance pointing to the same directory.
        let store2 = Store::open(&dir).unwrap();
        let signer = store2.signer(&info.id).unwrap();
        assert_eq!(signer.key_id(), info.id);

        // The reloaded signer must produce signatures verifiable with
        // the same public key.
        let verifier = {
            use crate::attestation::Verifier;
            use ed25519_dalek::VerifyingKey;
            let pk_bytes: [u8; 32] = info.public_key.try_into().unwrap();
            let vk = VerifyingKey::from_bytes(&pk_bytes).unwrap();
            let mut v = Verifier::new(std::collections::HashMap::new());
            v.add_key(info.id.clone(), vk);
            v
        };

        use crate::attestation::sign;
        use crate::statements::ActionStatement;
        let stmt   = ActionStatement::new("agent://test", "tool.call");
        let pt     = crate::statements::payload_type("action");
        let signed = sign(&pt, &stmt, signer.as_ref()).unwrap();
        verifier.verify(&signed.envelope).unwrap();

        cleanup(dir);
    }

    #[test]
    fn list_keys() {
        let (store, dir) = make_store();
        store.generate(true).unwrap();
        store.generate(false).unwrap();

        let keys = store.list().unwrap();
        assert_eq!(keys.len(), 2);
        assert_eq!(keys.iter().filter(|k| k.is_default).count(), 1);
        cleanup(dir);
    }

    #[test]
    fn no_default_key_errors() {
        let (store, dir) = make_store();
        assert!(store.default_signer().is_err());
        cleanup(dir);
    }
}