secretumvault/src/engines/kv.rs

use std::sync::Arc;

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;

use super::Engine;
use crate::core::SealMechanism;
use crate::crypto::{CryptoBackend, SymmetricAlgorithm};
use crate::error::{Result, VaultError};
use crate::storage::{EncryptedData, StorageBackend};

/// Individual version of a secret
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KVVersion {
    pub version: u64,
    pub data: Value,
    pub created_at: DateTime<Utc>,
    pub deleted: bool,
}

/// Secret with full version history
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KVSecret {
    pub path: String,
    pub versions: Vec<KVVersion>,
    pub current_version: u64,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
}

impl KVSecret {
    /// Create a new secret
    fn new(path: String) -> Self {
        let now = Utc::now();
        Self {
            path,
            versions: Vec::new(),
            current_version: 0,
            created_at: now,
            updated_at: now,
        }
    }

    /// Get the current version's data
    fn current_data(&self) -> Option<&Value> {
        self.versions
            .iter()
            .find(|v| v.version == self.current_version && !v.deleted)
            .map(|v| &v.data)
    }

    /// Add a new version
    fn add_version(&mut self, data: Value) {
        let new_version = self.current_version + 1;
        self.versions.push(KVVersion {
            version: new_version,
            data,
            created_at: Utc::now(),
            deleted: false,
        });
        self.current_version = new_version;
        self.updated_at = Utc::now();
    }

    /// Mark the current version as deleted (soft delete)
    fn soft_delete(&mut self) -> Result<()> {
        if let Some(version) = self
            .versions
            .iter_mut()
            .find(|v| v.version == self.current_version)
        {
            version.deleted = true;
            self.updated_at = Utc::now();
            Ok(())
        } else {
            Err(VaultError::storage("Version not found".to_string()))
        }
    }

    /// Get a specific version's data
    #[allow(dead_code)]
    fn get_version(&self, version: u64) -> Result<Option<&Value>> {
        Ok(self
            .versions
            .iter()
            .find(|v| v.version == version && !v.deleted)
            .map(|v| &v.data))
    }

    /// List all non-deleted versions
    #[allow(dead_code)]
    fn list_versions(&self) -> Vec<u64> {
        self.versions
            .iter()
            .filter(|v| !v.deleted)
            .map(|v| v.version)
            .collect()
    }
}

/// KV Secrets Engine (v2 with versioning)
#[derive(Debug)]
pub struct KVEngine {
    storage: Arc<dyn StorageBackend>,
    crypto: Arc<dyn CryptoBackend>,
    seal: Arc<tokio::sync::Mutex<SealMechanism>>,
    mount_path: String,
}

impl KVEngine {
    /// Create a new KV engine instance
    pub fn new(
        storage: Arc<dyn StorageBackend>,
        crypto: Arc<dyn CryptoBackend>,
        seal: Arc<tokio::sync::Mutex<SealMechanism>>,
        mount_path: String,
    ) -> Self {
        Self {
            storage,
            crypto,
            seal,
            mount_path,
        }
    }

    /// Get the storage key for a secret path
    fn storage_key(&self, path: &str) -> String {
        format!("{}data/{}", self.mount_path, path)
    }

    /// Encrypt secret data using master key
    async fn encrypt_secret(&self, data: &[u8]) -> Result<EncryptedData> {
        let seal = self.seal.lock().await;
        let master_key = seal
            .master_key()
            .map_err(|e| VaultError::crypto(e.to_string()))?;

        let ciphertext = self
            .crypto
            .encrypt_symmetric(&master_key.key_data, data, SymmetricAlgorithm::Aes256Gcm)
            .await
            .map_err(|e| VaultError::crypto(e.to_string()))?;

        // Extract nonce (first 12 bytes) and actual ciphertext
        let nonce = ciphertext[..12].to_vec();
        let ct = ciphertext[12..].to_vec();

        Ok(EncryptedData {
            ciphertext: ct,
            nonce,
            algorithm: "AES-256-GCM".to_string(),
        })
    }

    /// Decrypt secret data using master key
    async fn decrypt_secret(&self, encrypted: &EncryptedData) -> Result<Vec<u8>> {
        let seal = self.seal.lock().await;
        let master_key = seal
            .master_key()
            .map_err(|e| VaultError::crypto(e.to_string()))?;

        let mut combined = encrypted.nonce.clone();
        combined.extend_from_slice(&encrypted.ciphertext);

        self.crypto
            .decrypt_symmetric(
                &master_key.key_data,
                &combined,
                SymmetricAlgorithm::Aes256Gcm,
            )
            .await
            .map_err(|e| VaultError::crypto(e.to_string()))
    }

    /// Load secret from storage
    async fn load_secret(&self, path: &str) -> Result<Option<KVSecret>> {
        let key = self.storage_key(path);
        match self.storage.get_secret(&key).await {
            Ok(encrypted_data) => {
                let decrypted = self.decrypt_secret(&encrypted_data).await?;
                let secret: KVSecret = serde_json::from_slice(&decrypted)
                    .map_err(|e| VaultError::storage(e.to_string()))?;
                Ok(Some(secret))
            }
            Err(e) => {
                // Check if it's a NotFound error by examining the error message
                if e.to_string().contains("not found") || e.to_string().contains("Not found") {
                    Ok(None)
                } else {
                    Err(VaultError::storage(e.to_string()))
                }
            }
        }
    }

    /// Save secret to storage
    async fn save_secret(&self, secret: &KVSecret) -> Result<()> {
        let key = self.storage_key(&secret.path);
        let plaintext =
            serde_json::to_vec(secret).map_err(|e| VaultError::storage(e.to_string()))?;

        let encrypted = self.encrypt_secret(&plaintext).await?;

        self.storage
            .store_secret(&key, &encrypted)
            .await
            .map_err(|e| VaultError::storage(e.to_string()))
    }
}

#[async_trait]
impl Engine for KVEngine {
    fn name(&self) -> &str {
        "kv"
    }

    fn engine_type(&self) -> &str {
        "kv"
    }

    async fn read(&self, path: &str) -> Result<Option<Value>> {
        let secret = self.load_secret(path).await?;
        Ok(secret.and_then(|s| s.current_data().cloned()))
    }

    async fn write(&self, path: &str, data: &Value) -> Result<()> {
        let mut secret = match self.load_secret(path).await? {
            Some(s) => s,
            None => KVSecret::new(path.to_string()),
        };

        secret.add_version(data.clone());
        self.save_secret(&secret).await
    }

    async fn delete(&self, path: &str) -> Result<()> {
        let mut secret = self
            .load_secret(path)
            .await?
            .ok_or_else(|| VaultError::storage("Secret not found".to_string()))?;

        secret.soft_delete()?;
        self.save_secret(&secret).await
    }

    async fn list(&self, prefix: &str) -> Result<Vec<String>> {
        self.storage
            .list_secrets(&self.storage_key(prefix))
            .await
            .map_err(|e| VaultError::storage(e.to_string()))
    }

    async fn health_check(&self) -> Result<()> {
        self.storage
            .health_check()
            .await
            .map_err(|e| VaultError::storage(e.to_string()))?;

        let seal = self.seal.lock().await;
        if seal.is_sealed() {
            return Err(VaultError::crypto("Vault is sealed".to_string()));
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;
    use tempfile::TempDir;

    use super::*;
    use crate::config::{FilesystemStorageConfig, SealConfig, ShamirSealConfig, StorageConfig};
    use crate::crypto::CryptoRegistry;
    use crate::storage::StorageRegistry;

    async fn setup_engine() -> Result<(KVEngine, TempDir, Arc<dyn CryptoBackend>)> {
        let temp_dir = TempDir::new().map_err(|e| VaultError::storage(e.to_string()))?;

        let fs_config = FilesystemStorageConfig {
            path: temp_dir.path().to_path_buf(),
        };
        let storage_config = StorageConfig {
            backend: "filesystem".to_string(),
            filesystem: fs_config,
            surrealdb: Default::default(),
            etcd: Default::default(),
            postgresql: Default::default(),
        };

        let storage = StorageRegistry::create(&storage_config).await?;
        let crypto = CryptoRegistry::create("openssl", &Default::default())?;

        let seal_config = SealConfig {
            seal_type: "shamir".to_string(),
            shamir: ShamirSealConfig {
                threshold: 2,
                shares: 3,
            },
            auto_unseal: Default::default(),
        };
        let mut seal = SealMechanism::new(&seal_config)?;

        // Initialize and unseal for testing
        let _init_result = seal.init(crypto.as_ref(), storage.as_ref()).await?;

        let seal_arc = Arc::new(tokio::sync::Mutex::new(seal));

        let engine = KVEngine::new(storage, crypto.clone(), seal_arc, "secret/".to_string());

        Ok((engine, temp_dir, crypto))
    }

    #[tokio::test]
    async fn test_kv_write_and_read() -> Result<()> {
        let (engine, _temp, _) = setup_engine().await?;

        let data = json!({ "username": "admin", "password": "secret123" });
        engine.write("db/mysql", &data).await?;

        let read_data = engine.read("db/mysql").await?;
        assert_eq!(read_data, Some(data));

        Ok(())
    }

    #[tokio::test]
    async fn test_kv_versioning() -> Result<()> {
        let (engine, _temp, _) = setup_engine().await?;

        let data_v1 = json!({ "password": "old_password" });
        let data_v2 = json!({ "password": "new_password" });

        engine.write("app/api_key", &data_v1).await?;
        engine.write("app/api_key", &data_v2).await?;

        let current = engine.read("app/api_key").await?;
        assert_eq!(current, Some(data_v2));

        Ok(())
    }

    #[tokio::test]
    async fn test_kv_delete() -> Result<()> {
        let (engine, _temp, _) = setup_engine().await?;

        let data = json!({ "secret": "value" });
        engine.write("test/secret", &data).await?;

        assert!(engine.read("test/secret").await?.is_some());

        engine.delete("test/secret").await?;

        let deleted = engine.read("test/secret").await?;
        assert!(deleted.is_none());

        Ok(())
    }

    #[tokio::test]
    async fn test_kv_health_check() -> Result<()> {
        let (engine, _temp, _) = setup_engine().await?;
        engine.health_check().await?;
        Ok(())
    }
}