newton-core 0.4.16

//! Encrypted keystore for threshold key material.
//!
//! Provides AES-256-GCM-SIV encryption with scrypt KDF for persisting
//! threshold key shares (FROST KeyPackages and ThresholdDecryptionContexts).
//! Mirrors the EIP-2335 BLS keystore pattern used by Ethereum validators.
//!
//! # Security properties
//!
//! - scrypt (log_n=18, r=8, p=1) for memory-hard password stretching
//! - AES-256-GCM-SIV provides nonce-misuse resistance
//! - Random 32-byte salt and 12-byte nonce per write
//! - `ceremony_id` bound into stored metadata for provenance tracking
//!
//! # File format
//!
//! JSON with fields: `version`, `ceremony_id`, `created_at`, `encryption`.
//! The `encryption` sub-object holds KDF params and hex-encoded `salt`,
//! `nonce`, `ciphertext`.

use std::path::{Path, PathBuf};

use aes_gcm_siv::{
    aead::{Aead, KeyInit},
    Aes256GcmSiv, Key, Nonce,
};
use rand_core::{OsRng, RngCore};
use serde::{Deserialize, Serialize};
use zeroize::Zeroizing;

#[cfg(unix)]
use std::os::unix::fs::OpenOptionsExt;

use crate::crypto::error::CryptoError;

/// Version constant for the keystore file format.
const KEYSTORE_VERSION: u32 = 1;

/// scrypt parameters for key derivation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScryptParams {
    /// log2 of the work factor N (cost parameter).
    pub log_n: u8,
    /// Block size parameter.
    pub r: u32,
    /// Parallelization parameter.
    pub p: u32,
}

impl Default for ScryptParams {
    fn default() -> Self {
        Self { log_n: 18, r: 8, p: 1 }
    }
}

/// Encryption metadata stored alongside the ciphertext.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeystoreEncryption {
    /// KDF algorithm identifier.
    pub kdf: String,
    /// scrypt parameters.
    pub params: ScryptParams,
    /// Hex-encoded 32-byte random salt.
    pub salt: String,
    /// Hex-encoded 12-byte random nonce (AES-GCM-SIV).
    pub nonce: String,
    /// Hex-encoded AES-256-GCM-SIV ciphertext + tag.
    pub ciphertext: String,
}

/// On-disk keystore file structure.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ThresholdKeystore {
    /// Format version — must be 1.
    pub version: u32,
    /// DKG ceremony identifier this share belongs to.
    pub ceremony_id: String,
    /// RFC 3339 timestamp of when the keystore was written.
    pub created_at: String,
    /// Encryption envelope.
    pub encryption: KeystoreEncryption,
}

/// Derive a 32-byte AES key from `password` + `salt` using scrypt.
///
/// Returns a [`Zeroizing`] wrapper that automatically overwrites the key
/// material with zeros when dropped, preventing it from lingering on the stack.
fn derive_key(password: &[u8], salt: &[u8], params: &ScryptParams) -> Result<Zeroizing<[u8; 32]>, CryptoError> {
    let log_n = params.log_n;
    let r = params.r;
    let p = params.p;

    let scrypt_params = scrypt::Params::new(log_n, r, p, 32)
        .map_err(|e| CryptoError::KeystoreEncrypt(format!("invalid scrypt params: {e}")))?;

    let mut key = Zeroizing::new([0u8; 32]);
    scrypt::scrypt(password, salt, &scrypt_params, &mut *key)
        .map_err(|e| CryptoError::KeystoreEncrypt(format!("scrypt KDF failed: {e}")))?;
    Ok(key)
}

/// Encrypt `plaintext` with AES-256-GCM-SIV using the given `key` and `nonce`.
fn aes_encrypt(key: &[u8; 32], nonce_bytes: &[u8; 12], plaintext: &[u8]) -> Result<Vec<u8>, CryptoError> {
    let cipher = Aes256GcmSiv::new(Key::<Aes256GcmSiv>::from_slice(key));
    let nonce = Nonce::from_slice(nonce_bytes);
    cipher
        .encrypt(nonce, plaintext)
        .map_err(|e| CryptoError::KeystoreEncrypt(format!("AES-GCM-SIV encrypt failed: {e}")))
}

/// Decrypt `ciphertext` with AES-256-GCM-SIV using the given `key` and `nonce`.
fn aes_decrypt(key: &[u8; 32], nonce_bytes: &[u8; 12], ciphertext: &[u8]) -> Result<Vec<u8>, CryptoError> {
    let cipher = Aes256GcmSiv::new(Key::<Aes256GcmSiv>::from_slice(key));
    let nonce = Nonce::from_slice(nonce_bytes);
    cipher.decrypt(nonce, ciphertext).map_err(|_| {
        CryptoError::KeystoreDecrypt("AES-GCM-SIV decryption failed (wrong password or corrupted data)".into())
    })
}

/// Write an encrypted keystore to `path`.
///
/// Generates a fresh random salt and nonce on every call. The `plaintext`
/// is typically a serialized [`frost_ristretto255::keys::KeyPackage`] or
/// [`crate::dkg::types::ThresholdDecryptionContext`].
///
/// # Errors
///
/// Returns [`CryptoError::KeystoreEncrypt`] if scrypt or AES-GCM-SIV fail,
/// or if the file cannot be written.
pub fn write_keystore(path: &Path, ceremony_id: &str, plaintext: &[u8], password: &[u8]) -> Result<(), CryptoError> {
    let mut salt = [0u8; 32];
    OsRng.fill_bytes(&mut salt);

    let mut nonce_bytes = [0u8; 12];
    OsRng.fill_bytes(&mut nonce_bytes);

    let params = ScryptParams::default();
    let key = derive_key(password, &salt, &params)?;
    let ciphertext = aes_encrypt(&key, &nonce_bytes, plaintext)?;

    let salt_hex = hex::encode(salt);
    let nonce_hex = hex::encode(nonce_bytes);
    let ciphertext_hex = hex::encode(&ciphertext);
    let created_at = chrono::Utc::now().to_rfc3339();

    let keystore = ThresholdKeystore {
        version: KEYSTORE_VERSION,
        ceremony_id: ceremony_id.to_string(),
        created_at,
        encryption: KeystoreEncryption {
            kdf: "scrypt".to_string(),
            params,
            salt: salt_hex,
            nonce: nonce_hex,
            ciphertext: ciphertext_hex,
        },
    };

    let json = serde_json::to_string_pretty(&keystore)
        .map_err(|e| CryptoError::KeystoreEncrypt(format!("JSON serialization failed: {e}")))?;

    {
        use std::{fs::OpenOptions, io::Write};

        let mut opts = OpenOptions::new();
        opts.write(true).create(true).truncate(true);

        // Restrict to owner-only (0600) on Unix — keystore contains encrypted key material
        #[cfg(unix)]
        opts.mode(0o600);

        let mut file = opts
            .open(path)
            .map_err(|e| CryptoError::KeystoreEncrypt(format!("failed to open keystore at {}: {e}", path.display())))?;

        file.write_all(json.as_bytes()).map_err(|e| {
            CryptoError::KeystoreEncrypt(format!("failed to write keystore to {}: {e}", path.display()))
        })?;
    }

    Ok(())
}

/// Read and decrypt a keystore from `path`.
///
/// # Errors
///
/// Returns [`CryptoError::KeystoreDecrypt`] if the file is missing, the JSON
/// is malformed, the version is unsupported, the hex encoding is invalid, or
/// the AES-GCM-SIV decryption fails (wrong password or corrupted data).
pub fn read_keystore(path: &Path, password: &[u8]) -> Result<Vec<u8>, CryptoError> {
    let json = std::fs::read_to_string(path)
        .map_err(|e| CryptoError::KeystoreDecrypt(format!("failed to read keystore from {}: {e}", path.display())))?;

    let keystore: ThresholdKeystore =
        serde_json::from_str(&json).map_err(|e| CryptoError::KeystoreDecrypt(format!("invalid keystore JSON: {e}")))?;

    decrypt_keystore(&keystore, password)
}

/// Decrypt a parsed keystore.
pub fn decrypt_keystore(keystore: &ThresholdKeystore, password: &[u8]) -> Result<Vec<u8>, CryptoError> {
    if keystore.version != KEYSTORE_VERSION {
        return Err(CryptoError::KeystoreDecrypt(format!(
            "unsupported keystore version {} (expected {})",
            keystore.version, KEYSTORE_VERSION
        )));
    }

    if keystore.encryption.kdf != "scrypt" {
        return Err(CryptoError::KeystoreDecrypt(format!(
            "unsupported KDF: {} (expected scrypt)",
            keystore.encryption.kdf
        )));
    }

    let salt = hex::decode(&keystore.encryption.salt)
        .map_err(|e| CryptoError::KeystoreDecrypt(format!("invalid salt hex: {e}")))?;
    let salt: [u8; 32] = salt
        .try_into()
        .map_err(|_| CryptoError::KeystoreDecrypt("salt must be 32 bytes".into()))?;

    let nonce_bytes = hex::decode(&keystore.encryption.nonce)
        .map_err(|e| CryptoError::KeystoreDecrypt(format!("invalid nonce hex: {e}")))?;
    let nonce_bytes: [u8; 12] = nonce_bytes
        .try_into()
        .map_err(|_| CryptoError::KeystoreDecrypt("nonce must be 12 bytes".into()))?;

    let ciphertext = hex::decode(&keystore.encryption.ciphertext)
        .map_err(|e| CryptoError::KeystoreDecrypt(format!("invalid ciphertext hex: {e}")))?;

    let key = derive_key(password, &salt, &keystore.encryption.params)
        .map_err(|e| CryptoError::KeystoreDecrypt(format!("KDF failed: {e}")))?;

    aes_decrypt(&key, &nonce_bytes, &ciphertext)
}

/// Filename prefix for epoch-specific keystore files.
const EPOCH_KEYSTORE_PREFIX: &str = "threshold_keystore_epoch_";

/// Write an encrypted keystore for a specific epoch.
///
/// The file is written to `dir/threshold_keystore_epoch_{epoch_id}.json`.
/// Each epoch produces an independent keystore so operators can hold
/// key material for the current and previous epochs during grace periods.
///
/// # Errors
///
/// Returns [`CryptoError::KeystoreEncrypt`] on encryption or I/O failure.
pub fn write_epoch_keystore(
    dir: &Path,
    epoch_id: u64,
    ceremony_id: &str,
    plaintext: &[u8],
    password: &[u8],
) -> Result<PathBuf, CryptoError> {
    let filename = format!("{EPOCH_KEYSTORE_PREFIX}{epoch_id}.json");
    let path = dir.join(filename);
    write_keystore(&path, ceremony_id, plaintext, password)?;
    Ok(path)
}

/// List all epoch keystore files in `dir`, sorted by epoch ID descending.
///
/// Scans for files matching `threshold_keystore_epoch_{N}.json` and returns
/// `(epoch_id, path)` pairs. Non-matching files are silently ignored.
///
/// # Errors
///
/// Returns [`CryptoError::KeystoreDecrypt`] if the directory cannot be read.
pub fn list_epoch_keystores(dir: &Path) -> Result<Vec<(u64, PathBuf)>, CryptoError> {
    let entries = std::fs::read_dir(dir).map_err(|e| {
        CryptoError::KeystoreDecrypt(format!("failed to read keystore directory {}: {e}", dir.display()))
    })?;

    let mut keystores: Vec<(u64, PathBuf)> = Vec::new();

    for entry in entries {
        let entry = entry.map_err(|e| {
            CryptoError::KeystoreDecrypt(format!("failed to read directory entry in {}: {e}", dir.display()))
        })?;

        let path = entry.path();
        if let Some(epoch_id) = parse_epoch_id_from_path(&path) {
            keystores.push((epoch_id, path));
        }
    }

    // Sort descending by epoch ID so the latest epoch comes first
    keystores.sort_by_key(|k| std::cmp::Reverse(k.0));
    Ok(keystores)
}

/// Read and decrypt the keystore with the highest epoch ID in `dir`.
///
/// Returns `Ok(None)` if no epoch keystore files exist.
///
/// # Errors
///
/// Returns [`CryptoError::KeystoreDecrypt`] if the directory cannot be read
/// or the highest-epoch keystore cannot be decrypted.
pub fn read_latest_epoch_keystore(dir: &Path, password: &[u8]) -> Result<Option<(u64, Vec<u8>)>, CryptoError> {
    let keystores = list_epoch_keystores(dir)?;

    match keystores.first() {
        Some((epoch_id, path)) => {
            let plaintext = read_keystore(path, password)?;
            Ok(Some((*epoch_id, plaintext)))
        }
        None => Ok(None),
    }
}

/// Delete the epoch keystore file for `epoch_id` in `dir`.
///
/// Returns `Ok(())` if the file was removed or did not exist (idempotent).
///
/// # Errors
///
/// Returns [`CryptoError::KeystoreDecrypt`] if the file exists but cannot be removed.
pub fn delete_epoch_keystore(dir: &Path, epoch_id: u64) -> Result<(), CryptoError> {
    let filename = format!("{EPOCH_KEYSTORE_PREFIX}{epoch_id}.json");
    let path = dir.join(filename);
    match std::fs::remove_file(&path) {
        Ok(()) => Ok(()),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
        Err(e) => Err(CryptoError::KeystoreDecrypt(format!(
            "failed to delete epoch keystore at {}: {e}",
            path.display()
        ))),
    }
}

/// Extract the epoch ID from a keystore filename, if it matches the expected pattern.
fn parse_epoch_id_from_path(path: &Path) -> Option<u64> {
    let filename = path.file_name()?.to_str()?;
    let stem = filename.strip_prefix(EPOCH_KEYSTORE_PREFIX)?;
    let epoch_str = stem.strip_suffix(".json")?;
    epoch_str.parse::<u64>().ok()
}

#[cfg(test)]
mod tests {
    use tempfile::NamedTempFile;

    use super::*;

    // -- Epoch keystore tests --

    #[test]
    fn epoch_keystore_write_read_roundtrip() {
        let dir = tempfile::tempdir().expect("tempdir");
        let plaintext = b"epoch-0-key-share-data";
        let password = b"test-password";

        let path =
            write_epoch_keystore(dir.path(), 0, "ceremony-epoch-0", plaintext, password).expect("write_epoch_keystore");
        assert!(path.exists());

        let result = read_latest_epoch_keystore(dir.path(), password).expect("read_latest");
        let (epoch_id, recovered) = result.expect("expected Some");
        assert_eq!(epoch_id, 0);
        assert_eq!(recovered, plaintext);
    }

    #[test]
    fn list_epoch_keystores_sorted_descending() {
        let dir = tempfile::tempdir().expect("tempdir");
        let password = b"test-password";

        // Write epochs out of order: 0, 2, 1
        write_epoch_keystore(dir.path(), 0, "c0", b"data-0", password).expect("write 0");
        write_epoch_keystore(dir.path(), 2, "c2", b"data-2", password).expect("write 2");
        write_epoch_keystore(dir.path(), 1, "c1", b"data-1", password).expect("write 1");

        let keystores = list_epoch_keystores(dir.path()).expect("list");
        let epoch_ids: Vec<u64> = keystores.iter().map(|(id, _)| *id).collect();
        assert_eq!(epoch_ids, vec![2, 1, 0]);
    }

    #[test]
    fn delete_epoch_keystore_removes_file() {
        let dir = tempfile::tempdir().expect("tempdir");
        let password = b"test-password";

        write_epoch_keystore(dir.path(), 5, "c5", b"data-5", password).expect("write 5");
        let keystores = list_epoch_keystores(dir.path()).expect("list before delete");
        assert_eq!(keystores.len(), 1);

        delete_epoch_keystore(dir.path(), 5).expect("delete");
        let keystores = list_epoch_keystores(dir.path()).expect("list after delete");
        assert!(keystores.is_empty());
    }

    #[test]
    fn read_latest_returns_none_for_empty_dir() {
        let dir = tempfile::tempdir().expect("tempdir");
        let result = read_latest_epoch_keystore(dir.path(), b"any-password").expect("read_latest");
        assert!(result.is_none());
    }

    #[test]
    fn keystore_write_read_roundtrip() {
        let plaintext = b"this is my secret threshold key share data";
        let password = b"correct-horse-battery-staple";
        let ceremony_id = "test-ceremony-abc123";

        let tmp = NamedTempFile::new().expect("tempfile");
        write_keystore(tmp.path(), ceremony_id, plaintext, password).expect("write_keystore");
        let recovered = read_keystore(tmp.path(), password).expect("read_keystore");

        assert_eq!(recovered, plaintext);
    }

    #[test]
    fn keystore_wrong_password_fails() {
        let plaintext = b"sensitive key material";
        let password = b"correct-password";
        let wrong_password = b"wrong-password";

        let tmp = NamedTempFile::new().expect("tempfile");
        write_keystore(tmp.path(), "ceremony-1", plaintext, password).expect("write_keystore");

        let result = read_keystore(tmp.path(), wrong_password);
        assert!(result.is_err(), "expected decryption to fail with wrong password");
        let err = result.unwrap_err();
        assert!(
            matches!(err, CryptoError::KeystoreDecrypt(_)),
            "expected KeystoreDecrypt, got: {err}"
        );
    }

    #[test]
    fn keystore_file_not_found() {
        let path = Path::new("/tmp/nonexistent-newton-keystore-abc123.json");
        let result = read_keystore(path, b"password");
        assert!(result.is_err(), "expected error for nonexistent file");
        let err = result.unwrap_err();
        assert!(
            matches!(err, CryptoError::KeystoreDecrypt(_)),
            "expected KeystoreDecrypt, got: {err}"
        );
    }

    #[test]
    fn keystore_corrupted_payload() {
        let plaintext = b"key share bytes";
        let password = b"my-password";

        let tmp = NamedTempFile::new().expect("tempfile");
        write_keystore(tmp.path(), "ceremony-2", plaintext, password).expect("write_keystore");

        // Read the file, corrupt the ciphertext hex, write back
        let json = std::fs::read_to_string(tmp.path()).expect("read");
        let mut keystore: ThresholdKeystore = serde_json::from_str(&json).expect("parse");

        // Flip a byte in the ciphertext hex
        let mut ciphertext_bytes = hex::decode(&keystore.encryption.ciphertext).expect("decode");
        ciphertext_bytes[0] ^= 0xff;
        keystore.encryption.ciphertext = hex::encode(&ciphertext_bytes);

        let corrupted_json = serde_json::to_string_pretty(&keystore).expect("serialize");
        std::fs::write(tmp.path(), corrupted_json).expect("write");

        let result = read_keystore(tmp.path(), password);
        assert!(result.is_err(), "expected decryption to fail with corrupted payload");
        let err = result.unwrap_err();
        assert!(
            matches!(err, CryptoError::KeystoreDecrypt(_)),
            "expected KeystoreDecrypt, got: {err}"
        );
    }
}