foxhunt/services/ml_training_service/src/encryption.rs

//! Model Encryption Key Management with Vault Integration
//!
//! This module provides secure encryption key management for ML model storage,
//! supporting key rotation, multiple algorithms, and secure key retrieval from Vault.

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use tokio::fs;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};

use config::ConfigLoader;
use config::EncryptionConfig;

/// Encryption keys structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptionKeys {
    pub primary_key: String,
    pub key_id: String,
    pub algorithm: String,
    pub created_at: SystemTime,
}

impl EncryptionKeys {
    /// Check if key should be rotated based on age
    pub fn should_rotate(&self, rotation_days: u64) -> bool {
        match self.created_at.elapsed() {
            Ok(elapsed) => elapsed.as_secs() > rotation_days * 86400,
            Err(_) => true, // If we can't determine age, assume rotation needed
        }
    }
}

/// Encryption key manager with secure configuration
pub struct EncryptionKeyManager {
    config: EncryptionConfig,
    config_loader: Option<ConfigLoader>,
    cached_keys: Arc<RwLock<Option<CachedEncryptionKeys>>>,
}

/// Cached encryption keys with metadata
#[derive(Debug, Clone)]
struct CachedEncryptionKeys {
    keys: EncryptionKeys,
    cached_at: SystemTime,
    cache_ttl_secs: u64,
}

impl CachedEncryptionKeys {
    fn new(keys: EncryptionKeys, cache_ttl_secs: u64) -> Self {
        Self {
            keys,
            cached_at: SystemTime::now(),
            cache_ttl_secs,
        }
    }

    fn is_expired(&self) -> bool {
        match self.cached_at.elapsed() {
            Ok(elapsed) => elapsed.as_secs() > self.cache_ttl_secs,
            Err(_) => true, // If we can't determine time, assume expired
        }
    }

    fn needs_rotation(&self, rotation_days: u64) -> bool {
        self.keys.should_rotate(rotation_days)
    }
}

/// Encryption algorithm configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptionAlgorithmConfig {
    pub algorithm: EncryptionAlgorithm,
    pub key_size_bits: usize,
    pub iv_size_bytes: usize,
    pub tag_size_bytes: usize,
}

/// Supported encryption algorithms
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum EncryptionAlgorithm {
    #[serde(rename = "AES-256-GCM")]
    Aes256Gcm,
    #[serde(rename = "ChaCha20Poly1305")]
    ChaCha20Poly1305,
    #[serde(rename = "AES-256-CTR")]
    Aes256Ctr,
}

impl EncryptionAlgorithm {
    /// Get algorithm configuration
    pub fn get_config(&self) -> EncryptionAlgorithmConfig {
        match self {
            Self::Aes256Gcm => EncryptionAlgorithmConfig {
                algorithm: self.clone(),
                key_size_bits: 256,
                iv_size_bytes: 12,  // 96-bit IV for GCM
                tag_size_bytes: 16, // 128-bit authentication tag
            },
            Self::ChaCha20Poly1305 => EncryptionAlgorithmConfig {
                algorithm: self.clone(),
                key_size_bits: 256,
                iv_size_bytes: 12,  // 96-bit nonce
                tag_size_bytes: 16, // 128-bit authentication tag
            },
            Self::Aes256Ctr => EncryptionAlgorithmConfig {
                algorithm: self.clone(),
                key_size_bits: 256,
                iv_size_bytes: 16, // 128-bit IV for CTR
                tag_size_bytes: 0, // No authentication tag (not AEAD)
            },
        }
    }

    /// Check if algorithm provides authenticated encryption
    pub fn is_authenticated(&self) -> bool {
        matches!(self, Self::Aes256Gcm | Self::ChaCha20Poly1305)
    }
}

impl std::str::FromStr for EncryptionAlgorithm {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self> {
        match s {
            "AES-256-GCM" => Ok(Self::Aes256Gcm),
            "ChaCha20Poly1305" => Ok(Self::ChaCha20Poly1305),
            "AES-256-CTR" => Ok(Self::Aes256Ctr),
            _ => Err(anyhow::anyhow!("Unsupported encryption algorithm: {}", s)),
        }
    }
}

impl std::fmt::Display for EncryptionAlgorithm {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let name = match self {
            Self::Aes256Gcm => "AES-256-GCM",
            Self::ChaCha20Poly1305 => "ChaCha20Poly1305",
            Self::Aes256Ctr => "AES-256-CTR",
        };
        write!(f, "{}", name)
    }
}

/// Encryption metadata for stored models
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptionMetadata {
    pub algorithm: EncryptionAlgorithm,
    pub key_id: String,
    pub iv: Vec<u8>,
    pub tag: Option<Vec<u8>>, // For AEAD algorithms
    pub encrypted_at: SystemTime,
    pub key_version: u32,
}

impl EncryptionKeyManager {
    /// Create a new encryption key manager
    pub fn new(config: EncryptionConfig, config_loader: Option<ConfigLoader>) -> Self {
        Self {
            config,
            config_loader,
            cached_keys: Arc::new(RwLock::new(None)),
        }
    }

    /// Check if encryption is enabled
    pub fn is_encryption_enabled(&self) -> bool {
        self.config.enable_encryption
    }

    /// Get current encryption algorithm
    pub fn get_algorithm(&self) -> Result<EncryptionAlgorithm> {
        self.config.algorithm.parse()
    }

    /// Load encryption keys from secure configuration or fallback source
    pub async fn load_encryption_keys(&self) -> Result<EncryptionKeys> {
        // Check cache first
        {
            let cached_guard = self.cached_keys.read().await;
            if let Some(cached) = cached_guard.as_ref() {
                if !cached.is_expired() {
                    debug!("Using cached encryption keys");
                    return Ok(cached.keys.clone());
                }
            }
        }

        // Try to load from secure configuration first
        let keys = if let Some(config_loader) = &self.config_loader {
            match config_loader.get_encryption_keys().await {
                Ok(keys) => {
                    info!("Successfully loaded encryption keys from secure configuration");
                    keys
                }
                Err(e) => {
                    warn!("Failed to load encryption keys from secure configuration, trying fallback: {}", e);
                    self.load_fallback_keys().await?
                }
            }
        } else {
            info!("Loading encryption keys from fallback source (no secure configuration)");
            self.load_fallback_keys().await?
        };

        // Cache the loaded keys
        {
            let mut cached_guard = self.cached_keys.write().await;
            *cached_guard = Some(CachedEncryptionKeys::new(keys.clone(), 300)); // 5-minute cache
        }

        debug!("Encryption keys loaded and cached");
        Ok(keys)
    }

    /// Load encryption keys from fallback source (local file or generated)
    async fn load_fallback_keys(&self) -> Result<EncryptionKeys> {
        if let Some(key_file) = &self.config.local_key_file {
            self.load_keys_from_file(key_file).await
        } else {
            warn!("No fallback key source configured, generating temporary keys");
            self.generate_temporary_keys().await
        }
    }

    /// Load encryption keys from local file
    async fn load_keys_from_file(&self, key_file: &PathBuf) -> Result<EncryptionKeys> {
        let key_data = fs::read_to_string(key_file)
            .await
            .context("Failed to read encryption key file")?;

        let keys: EncryptionKeys =
            serde_json::from_str(&key_data).context("Failed to parse encryption key file")?;

        info!("Loaded encryption keys from file: {}", key_file.display());
        Ok(keys)
    }

    /// Generate temporary encryption keys (for development/fallback)
    async fn generate_temporary_keys(&self) -> Result<EncryptionKeys> {
        warn!("Generating temporary encryption keys - NOT suitable for production!");

        // Generate a random key (in production, use proper cryptographic libraries)
        let key_bytes: Vec<u8> = (0..32).map(|_| rand::random::<u8>()).collect();
        let primary_key = base64::encode(&key_bytes);

        let keys = EncryptionKeys {
            primary_key,
            key_id: format!(
                "temp-key-{}",
                SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs()
            ),
            algorithm: self.config.algorithm.clone(),
            created_at: SystemTime::now(),
        };

        Ok(keys)
    }

    /// Check if current keys need rotation
    pub async fn keys_need_rotation(&self) -> Result<bool> {
        let keys = self.load_encryption_keys().await?;
        Ok(keys.should_rotate(self.config.key_rotation_days))
    }

    /// Request key rotation (would trigger Vault key rotation in production)
    pub async fn request_key_rotation(&self) -> Result<()> {
        info!("Requesting encryption key rotation");

        // Clear cache to force reload of new keys
        {
            let mut cached_guard = self.cached_keys.write().await;
            *cached_guard = None;
        }

        // In a production environment, this would:
        // 1. Request new keys from Vault
        // 2. Update Vault secrets with new keys
        // 3. Notify other services of key rotation
        // 4. Schedule old key deprecation

        warn!("Key rotation requested - implement Vault key rotation logic for production");
        Ok(())
    }

    /// Encrypt model data
    pub async fn encrypt_model_data(&self, data: &[u8]) -> Result<(Vec<u8>, EncryptionMetadata)> {
        if !self.config.enable_encryption {
            return Err(anyhow::anyhow!("Encryption is disabled"));
        }

        let keys = self.load_encryption_keys().await?;
        let algorithm = self.get_algorithm()?;
        let algo_config = algorithm.get_config();

        // Generate random IV/nonce
        let iv: Vec<u8> = (0..algo_config.iv_size_bytes)
            .map(|_| rand::random::<u8>())
            .collect();

        // In production, use proper cryptographic libraries like ring, rustcrypto, etc.
        // This is a simplified implementation for demonstration
        let encrypted_data = self.perform_encryption(data, &keys.primary_key, &iv, &algorithm)?;

        let metadata = EncryptionMetadata {
            algorithm,
            key_id: keys.key_id.clone(),
            iv,
            tag: None, // Would be populated by actual AEAD encryption
            encrypted_at: SystemTime::now(),
            key_version: 1, // Would track actual key versions
        };

        debug!(
            "Encrypted {} bytes of model data with algorithm {}",
            data.len(),
            metadata.algorithm
        );

        Ok((encrypted_data, metadata))
    }

    /// Decrypt model data
    pub async fn decrypt_model_data(
        &self,
        encrypted_data: &[u8],
        metadata: &EncryptionMetadata,
    ) -> Result<Vec<u8>> {
        if !self.config.enable_encryption {
            return Err(anyhow::anyhow!("Encryption is disabled"));
        }

        let keys = self.load_encryption_keys().await?;

        // Verify key ID matches (in production, support multiple key versions)
        if keys.key_id != metadata.key_id {
            return Err(anyhow::anyhow!(
                "Key ID mismatch: expected {}, got {}",
                keys.key_id,
                metadata.key_id
            ));
        }

        let decrypted_data = self.perform_decryption(
            encrypted_data,
            &keys.primary_key,
            &metadata.iv,
            &metadata.algorithm,
        )?;

        debug!(
            "Decrypted {} bytes of model data with algorithm {}",
            decrypted_data.len(),
            metadata.algorithm
        );

        Ok(decrypted_data)
    }

    /// Perform actual encryption (simplified implementation)
    fn perform_encryption(
        &self,
        data: &[u8],
        key: &str,
        iv: &[u8],
        algorithm: &EncryptionAlgorithm,
    ) -> Result<Vec<u8>> {
        // This is a placeholder implementation
        // In production, use proper cryptographic libraries
        match algorithm {
            EncryptionAlgorithm::Aes256Gcm => {
                // Use AES-256-GCM encryption
                self.aes_gcm_encrypt(data, key, iv)
            }
            EncryptionAlgorithm::ChaCha20Poly1305 => {
                // Use ChaCha20-Poly1305 encryption
                self.chacha20_encrypt(data, key, iv)
            }
            EncryptionAlgorithm::Aes256Ctr => {
                // Use AES-256-CTR encryption
                self.aes_ctr_encrypt(data, key, iv)
            }
        }
    }

    /// Perform actual decryption (simplified implementation)
    fn perform_decryption(
        &self,
        encrypted_data: &[u8],
        key: &str,
        iv: &[u8],
        algorithm: &EncryptionAlgorithm,
    ) -> Result<Vec<u8>> {
        // This is a placeholder implementation
        // In production, use proper cryptographic libraries
        match algorithm {
            EncryptionAlgorithm::Aes256Gcm => {
                // Use AES-256-GCM decryption
                self.aes_gcm_decrypt(encrypted_data, key, iv)
            }
            EncryptionAlgorithm::ChaCha20Poly1305 => {
                // Use ChaCha20-Poly1305 decryption
                self.chacha20_decrypt(encrypted_data, key, iv)
            }
            EncryptionAlgorithm::Aes256Ctr => {
                // Use AES-256-CTR decryption
                self.aes_ctr_decrypt(encrypted_data, key, iv)
            }
        }
    }

    // Placeholder encryption methods (use proper crypto libraries in production)
    fn aes_gcm_encrypt(&self, data: &[u8], _key: &str, _iv: &[u8]) -> Result<Vec<u8>> {
        // Placeholder: XOR with pattern (NOT secure)
        warn!("Using placeholder AES-GCM encryption - implement proper crypto for production");
        Ok(data
            .iter()
            .enumerate()
            .map(|(i, &b)| b ^ ((i % 256) as u8))
            .collect())
    }

    fn aes_gcm_decrypt(&self, encrypted_data: &[u8], _key: &str, _iv: &[u8]) -> Result<Vec<u8>> {
        // Placeholder: XOR with pattern (NOT secure)
        warn!("Using placeholder AES-GCM decryption - implement proper crypto for production");
        Ok(encrypted_data
            .iter()
            .enumerate()
            .map(|(i, &b)| b ^ ((i % 256) as u8))
            .collect())
    }

    fn chacha20_encrypt(&self, data: &[u8], _key: &str, _iv: &[u8]) -> Result<Vec<u8>> {
        // Placeholder: Simple rotation (NOT secure)
        warn!("Using placeholder ChaCha20 encryption - implement proper crypto for production");
        Ok(data.iter().map(|&b| b.wrapping_add(1)).collect())
    }

    fn chacha20_decrypt(&self, encrypted_data: &[u8], _key: &str, _iv: &[u8]) -> Result<Vec<u8>> {
        // Placeholder: Simple rotation (NOT secure)
        warn!("Using placeholder ChaCha20 decryption - implement proper crypto for production");
        Ok(encrypted_data.iter().map(|&b| b.wrapping_sub(1)).collect())
    }

    fn aes_ctr_encrypt(&self, data: &[u8], _key: &str, _iv: &[u8]) -> Result<Vec<u8>> {
        // Placeholder: Byte reversal (NOT secure)
        warn!("Using placeholder AES-CTR encryption - implement proper crypto for production");
        Ok(data.iter().rev().cloned().collect())
    }

    fn aes_ctr_decrypt(&self, encrypted_data: &[u8], _key: &str, _iv: &[u8]) -> Result<Vec<u8>> {
        // Placeholder: Byte reversal (NOT secure)
        warn!("Using placeholder AES-CTR decryption - implement proper crypto for production");
        Ok(encrypted_data.iter().rev().cloned().collect())
    }

    /// Get encryption statistics
    pub async fn get_encryption_stats(&self) -> Result<EncryptionStats> {
        let keys = self.load_encryption_keys().await?;
        let needs_rotation = keys.should_rotate(self.config.key_rotation_days);

        Ok(EncryptionStats {
            encryption_enabled: self.config.enable_encryption,
            algorithm: self.get_algorithm()?,
            key_id: keys.key_id,
            key_age_days: keys
                .created_at
                .elapsed()
                .map(|d| d.as_secs() / 86400)
                .unwrap_or(0),
            needs_rotation,
            rotation_interval_days: self.config.key_rotation_days,
            cache_hit_rate: 0.0, // Would track actual cache statistics
        })
    }
}

/// Encryption statistics
#[derive(Debug, Clone, Serialize)]
pub struct EncryptionStats {
    pub encryption_enabled: bool,
    pub algorithm: EncryptionAlgorithm,
    pub key_id: String,
    pub key_age_days: u64,
    pub needs_rotation: bool,
    pub rotation_interval_days: u64,
    pub cache_hit_rate: f64,
}

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

    #[tokio::test]
    async fn test_encryption_algorithm_parsing() {
        assert_eq!(
            "AES-256-GCM".parse::<EncryptionAlgorithm>().unwrap(),
            EncryptionAlgorithm::Aes256Gcm
        );
        assert_eq!(
            "ChaCha20Poly1305".parse::<EncryptionAlgorithm>().unwrap(),
            EncryptionAlgorithm::ChaCha20Poly1305
        );
        assert!("Invalid-Algorithm".parse::<EncryptionAlgorithm>().is_err());
    }

    #[test]
    fn test_algorithm_config() {
        let aes_config = EncryptionAlgorithm::Aes256Gcm.get_config();
        assert_eq!(aes_config.key_size_bits, 256);
        assert_eq!(aes_config.iv_size_bytes, 12);
        assert!(aes_config.algorithm.is_authenticated());

        let ctr_config = EncryptionAlgorithm::Aes256Ctr.get_config();
        assert!(!ctr_config.algorithm.is_authenticated());
    }

    #[tokio::test]
    async fn test_encryption_key_manager_creation() {
        let config = EncryptionConfig {
            enable_encryption: true,
            algorithm: "AES-256-GCM".to_string(),
            key_rotation_days: 30,
            encryption_keys_vault_path: None,
            local_key_file: None,
        };

        let manager = EncryptionKeyManager::new(config, None);
        assert!(manager.is_encryption_enabled());
        assert_eq!(
            manager.get_algorithm().unwrap(),
            EncryptionAlgorithm::Aes256Gcm
        );
    }

    #[tokio::test]
    async fn test_temporary_key_generation() {
        let config = EncryptionConfig {
            enable_encryption: true,
            algorithm: "AES-256-GCM".to_string(),
            key_rotation_days: 30,
            encryption_keys_vault_path: None,
            local_key_file: None,
        };

        let manager = EncryptionKeyManager::new(config, None);
        let keys = manager.generate_temporary_keys().await.unwrap();

        assert!(!keys.primary_key.is_empty());
        assert!(!keys.key_id.is_empty());
        assert!(keys.key_id.starts_with("temp-key-"));
    }

    #[tokio::test]
    async fn test_placeholder_encryption_decryption() {
        let config = EncryptionConfig {
            enable_encryption: true,
            algorithm: "AES-256-GCM".to_string(),
            key_rotation_days: 30,
            encryption_keys_vault_path: None,
            local_key_file: None,
        };

        let manager = EncryptionKeyManager::new(config, None);
        let test_data = b"Hello, encrypted world!";

        let (encrypted, metadata) = manager.encrypt_model_data(test_data).await.unwrap();
        assert_ne!(encrypted, test_data);
        assert_eq!(metadata.algorithm, EncryptionAlgorithm::Aes256Gcm);

        let decrypted = manager
            .decrypt_model_data(&encrypted, &metadata)
            .await
            .unwrap();
        assert_eq!(decrypted, test_data);
    }
}