secretumvault/src/core/vault.rs

use std::collections::HashMap;
use std::sync::Arc;

use crate::auth::TokenManager;
#[cfg(feature = "server")]
use crate::background::LeaseRevocationWorker;
use crate::config::VaultConfig;
use crate::crypto::CryptoBackend;
use crate::engines::{DatabaseEngine, Engine, KVEngine, PkiEngine, TransitEngine};
use crate::error::Result;
use crate::storage::StorageBackend;
use crate::telemetry::Metrics;

/// Vault core - manages engines, crypto backend, and storage
pub struct VaultCore {
    /// Mounted secrets engines (mount_path -> engine)
    pub engines: HashMap<String, Box<dyn Engine>>,
    /// Storage backend
    pub storage: Arc<dyn StorageBackend>,
    /// Crypto backend
    pub crypto: Arc<dyn CryptoBackend>,
    /// Seal mechanism (behind a mutex for thread-safe unseal operations)
    pub seal: Arc<tokio::sync::Mutex<super::SealMechanism>>,
    /// Token manager for authentication and authorization
    pub token_manager: Arc<TokenManager>,
    /// Metrics collection
    pub metrics: Arc<Metrics>,
    /// Background lease revocation worker (server only)
    #[cfg(feature = "server")]
    pub lease_revocation_worker: Arc<LeaseRevocationWorker>,
}

impl VaultCore {
    /// Create vault core from configuration
    pub async fn from_config(config: &VaultConfig) -> Result<Self> {
        let storage = crate::storage::StorageRegistry::create(&config.storage).await?;
        let crypto =
            crate::crypto::CryptoRegistry::create(&config.vault.crypto_backend, &config.crypto)?;

        let seal_config = &config.seal;
        let seal = super::SealMechanism::new(seal_config)?;
        let seal = Arc::new(tokio::sync::Mutex::new(seal));

        let mut engines = HashMap::new();
        EngineRegistry::mount_engines(&config.engines, &storage, &crypto, &seal, &mut engines)?;

        // Initialize token manager with default 24-hour TTL
        let token_manager = Arc::new(TokenManager::new(storage.clone(), crypto.clone(), 24));

        // Initialize metrics
        let metrics = Arc::new(Metrics::new());

        #[cfg(feature = "server")]
        let lease_revocation_worker = {
            use crate::background::RevocationConfig;

            let revocation_config = RevocationConfig::default();
            let worker = Arc::new(LeaseRevocationWorker::new(
                storage.clone(),
                revocation_config,
            ));
            worker.start().await?;
            worker
        };

        Ok(Self {
            engines,
            storage,
            crypto,
            seal,
            token_manager,
            metrics,
            #[cfg(feature = "server")]
            lease_revocation_worker,
        })
    }

    /// Find engine by path prefix
    pub fn route_to_engine(&self, path: &str) -> Option<&dyn Engine> {
        let mut best_match: Option<(&str, &dyn Engine)> = None;

        for (mount_path, engine) in &self.engines {
            if !path.starts_with(mount_path) {
                continue;
            }

            let should_update = best_match
                .map(|(best_path, _)| mount_path.len() > best_path.len())
                .unwrap_or(true);

            if should_update {
                best_match = Some((mount_path, engine.as_ref()));
            }
        }

        best_match.map(|(_, engine)| engine)
    }

    /// Get the engine path and relative path after the mount point
    pub fn split_path(&self, path: &str) -> Option<(String, String)> {
        let mut best_match: Option<(&str, &str)> = None;

        for mount_path in self.engines.keys() {
            if !path.starts_with(mount_path) {
                continue;
            }

            let should_update = best_match
                .map(|(best_path, _)| mount_path.len() > best_path.len())
                .unwrap_or(true);

            if should_update {
                best_match = Some((mount_path, path));
            }
        }

        best_match.map(|(mount_path, path)| {
            let relative = if path.len() > mount_path.len() {
                path[mount_path.len()..].to_string()
            } else {
                String::new()
            };
            (mount_path.to_string(), relative)
        })
    }
}

/// Registry for creating and mounting engines
pub struct EngineRegistry;

impl EngineRegistry {
    /// Mount engines from configuration
    pub fn mount_engines(
        engines_config: &crate::config::EnginesConfig,
        storage: &Arc<dyn StorageBackend>,
        crypto: &Arc<dyn CryptoBackend>,
        seal: &Arc<tokio::sync::Mutex<super::SealMechanism>>,
        engines: &mut HashMap<String, Box<dyn Engine>>,
    ) -> Result<()> {
        // Mount KV engine from config
        if let Some(kv_config) = &engines_config.kv {
            let engine = KVEngine::new(
                storage.clone(),
                crypto.clone(),
                seal.clone(),
                kv_config.path.clone(),
            );
            engines.insert(kv_config.path.clone(), Box::new(engine) as Box<dyn Engine>);
        }

        // Mount Transit engine from config
        if let Some(transit_config) = &engines_config.transit {
            let engine = TransitEngine::new(
                storage.clone(),
                crypto.clone(),
                seal.clone(),
                transit_config.path.clone(),
            );
            engines.insert(
                transit_config.path.clone(),
                Box::new(engine) as Box<dyn Engine>,
            );
        }

        // Mount PKI engine from config
        if let Some(pki_config) = &engines_config.pki {
            let engine = PkiEngine::new(
                storage.clone(),
                crypto.clone(),
                seal.clone(),
                pki_config.path.clone(),
            );
            engines.insert(pki_config.path.clone(), Box::new(engine) as Box<dyn Engine>);
        }

        // Mount Database engine from config
        if let Some(database_config) = &engines_config.database {
            let engine = DatabaseEngine::new(
                storage.clone(),
                crypto.clone(),
                seal.clone(),
                database_config.path.clone(),
            );
            engines.insert(
                database_config.path.clone(),
                Box::new(engine) as Box<dyn Engine>,
            );
        }

        Ok(())
    }
}

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

    use super::*;
    use crate::config::{
        EngineConfig, FilesystemStorageConfig, SealConfig, ShamirSealConfig, StorageConfig,
    };

    fn create_test_vault_config(temp_dir: &TempDir) -> VaultConfig {
        VaultConfig {
            vault: Default::default(),
            server: Default::default(),
            storage: StorageConfig {
                backend: "filesystem".to_string(),
                filesystem: FilesystemStorageConfig {
                    path: temp_dir.path().to_path_buf(),
                },
                surrealdb: Default::default(),
                etcd: Default::default(),
                postgresql: Default::default(),
            },
            crypto: Default::default(),
            seal: SealConfig {
                seal_type: "shamir".to_string(),
                shamir: ShamirSealConfig {
                    threshold: 2,
                    shares: 3,
                },
                auto_unseal: Default::default(),
            },
            auth: Default::default(),
            engines: crate::config::EnginesConfig {
                kv: Some(EngineConfig {
                    path: "secret/".to_string(),
                    versioned: true,
                    extra: HashMap::new(),
                }),
                transit: None,
                pki: None,
                database: None,
            },
            logging: Default::default(),
            telemetry: Default::default(),
        }
    }

    #[tokio::test]
    async fn test_vault_core_creation() -> Result<()> {
        let temp_dir = TempDir::new().map_err(|e| crate::VaultError::storage(e.to_string()))?;
        let vault_config = create_test_vault_config(&temp_dir);

        let vault = VaultCore::from_config(&vault_config).await?;
        assert!(vault.engines.contains_key("secret/"));

        Ok(())
    }

    #[tokio::test]
    async fn test_route_to_engine() -> Result<()> {
        let temp_dir = TempDir::new().map_err(|e| crate::VaultError::storage(e.to_string()))?;
        let vault_config = create_test_vault_config(&temp_dir);

        let vault = VaultCore::from_config(&vault_config).await?;

        // Test routing
        let engine = vault.route_to_engine("secret/db/postgres");
        assert!(engine.is_some());
        assert_eq!(engine.unwrap().engine_type(), "kv");

        Ok(())
    }

    #[tokio::test]
    async fn test_split_path() -> Result<()> {
        let temp_dir = TempDir::new().map_err(|e| crate::VaultError::storage(e.to_string()))?;
        let vault_config = create_test_vault_config(&temp_dir);

        let vault = VaultCore::from_config(&vault_config).await?;

        let (mount_path, relative_path) = vault
            .split_path("secret/db/postgres")
            .expect("Failed to split path");

        assert_eq!(mount_path, "secret/");
        assert_eq!(relative_path, "db/postgres");

        Ok(())
    }

    #[tokio::test]
    async fn test_transit_engine_mounting() -> Result<()> {
        let temp_dir = TempDir::new().map_err(|e| crate::VaultError::storage(e.to_string()))?;

        let mut config = create_test_vault_config(&temp_dir);
        config.engines.transit = Some(EngineConfig {
            path: "transit/".to_string(),
            versioned: true,
            extra: HashMap::new(),
        });

        let vault = VaultCore::from_config(&config).await?;

        // Verify both KV and Transit engines are mounted
        assert!(vault.engines.contains_key("secret/"));
        assert!(vault.engines.contains_key("transit/"));

        // Verify routing to transit engine
        let engine = vault.route_to_engine("transit/keys/my-key");
        assert!(engine.is_some());
        assert_eq!(engine.unwrap().engine_type(), "transit");

        Ok(())
    }

    #[tokio::test]
    async fn test_pki_engine_mounting() -> Result<()> {
        let temp_dir = TempDir::new().map_err(|e| crate::VaultError::storage(e.to_string()))?;

        let mut config = create_test_vault_config(&temp_dir);
        config.engines.pki = Some(EngineConfig {
            path: "pki/".to_string(),
            versioned: false,
            extra: HashMap::new(),
        });

        let vault = VaultCore::from_config(&config).await?;

        // Verify both KV and PKI engines are mounted
        assert!(vault.engines.contains_key("secret/"));
        assert!(vault.engines.contains_key("pki/"));

        // Verify routing to PKI engine
        let engine = vault.route_to_engine("pki/certs/my-cert");
        assert!(engine.is_some());
        assert_eq!(engine.unwrap().engine_type(), "pki");

        Ok(())
    }

    #[tokio::test]
    async fn test_database_engine_mounting() -> Result<()> {
        let temp_dir = TempDir::new().map_err(|e| crate::VaultError::storage(e.to_string()))?;

        let mut config = create_test_vault_config(&temp_dir);
        config.engines.database = Some(EngineConfig {
            path: "database/".to_string(),
            versioned: false,
            extra: HashMap::new(),
        });

        let vault = VaultCore::from_config(&config).await?;

        // Verify both KV and Database engines are mounted
        assert!(vault.engines.contains_key("secret/"));
        assert!(vault.engines.contains_key("database/"));

        // Verify routing to database engine
        let engine = vault.route_to_engine("database/creds/postgres");
        assert!(engine.is_some());
        assert_eq!(engine.unwrap().engine_type(), "database");

        Ok(())
    }
}