reasonkit-core 0.1.8

//! Secure API Key Management
//!
//! Implements enterprise-grade API key generation, storage, and validation
//! with envelope encryption and constant-time comparison.
//!
//! # Security Features
//!
//! - **Envelope Encryption**: Keys encrypted with KEK from HSM/KMS
//! - **Hash-Only Storage**: Only SHA-256 hashes stored for validation
//! - **Constant-Time Comparison**: Prevents timing attacks
//! - **Zeroization**: Sensitive memory cleared after use
//! - **Scope-Based Access**: Fine-grained permission control
//!
//! # Key Format
//!
//! API keys follow the format: `rk_live_<base58_encoded_random>`
//! - Prefix `rk_live_` or `rk_test_` indicates environment
//! - Base58 encoding for URL-safe representation
//! - 32 bytes of cryptographically secure randomness

use crate::security::audit::{AuditEvent, AuditEventType, AuditLogger, AuditActor, AuditOutcome};
use chrono::{DateTime, Utc};
use secrecy::{ExposeSecret, SecretString, SecretVec};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use uuid::Uuid;

/// API Key metadata stored in database
///
/// IMPORTANT: This struct NEVER contains plaintext key material.
/// Only the hash and encrypted form are stored.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ApiKeyRecord {
    /// Unique identifier for this key
    pub key_id: Uuid,
    /// Tenant/organization this key belongs to
    pub tenant_id: Uuid,
    /// Key prefix for identification (first 16 chars only)
    pub key_prefix: String,
    /// SHA-256 hash of the full key (for validation)
    #[serde(with = "hex_array")]
    pub key_hash: [u8; 32],
    /// Encrypted key material (envelope encrypted)
    #[serde(with = "base64_vec")]
    pub encrypted_key: Vec<u8>,
    /// Nonce used for encryption
    #[serde(with = "hex_nonce")]
    pub encryption_nonce: [u8; 12],
    /// Key scope/permissions
    pub scopes: Vec<String>,
    /// Creation timestamp
    pub created_at: DateTime<Utc>,
    /// Expiration timestamp (optional)
    pub expires_at: Option<DateTime<Utc>>,
    /// Last used timestamp
    pub last_used_at: Option<DateTime<Utc>>,
    /// Key status
    pub status: KeyStatus,
    /// Rotation metadata
    pub rotation: Option<KeyRotationInfo>,
}

/// Key lifecycle status
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum KeyStatus {
    /// Key is active and valid
    Active,
    /// Key is being rotated (both old and new valid)
    Rotating,
    /// Key has been revoked
    Revoked,
    /// Key has expired
    Expired,
}

/// Rotation tracking metadata
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct KeyRotationInfo {
    /// Previous key ID (for rotation tracking)
    pub previous_key_id: Option<Uuid>,
    /// Rotation initiated at
    pub rotation_started: DateTime<Utc>,
    /// Rotation completed at
    pub rotation_completed: Option<DateTime<Utc>>,
    /// Rotation reason
    pub reason: RotationReason,
}

/// Reasons for key rotation
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RotationReason {
    /// Scheduled rotation (policy-based)
    Scheduled,
    /// Security incident detected
    SecurityIncident,
    /// Policy violation detected
    PolicyViolation,
    /// Customer-requested rotation
    CustomerRequest,
    /// Key potentially compromised
    KeyCompromise,
}

/// Generated key result (returned only once)
pub struct GeneratedKey {
    /// Key identifier
    pub key_id: Uuid,
    /// The plaintext API key (only returned once, never stored)
    pub api_key: SecretString,
    /// Expiration timestamp
    pub expires_at: Option<DateTime<Utc>>,
}

/// Validated key information
#[derive(Debug, Clone)]
pub struct ValidatedKey {
    /// Key identifier
    pub key_id: Uuid,
    /// Tenant identifier
    pub tenant_id: Uuid,
    /// Granted scopes
    pub scopes: Vec<String>,
}

/// Key validation errors
#[derive(Debug, thiserror::Error)]
pub enum KeyValidationError {
    #[error("Invalid API key")]
    InvalidKey,
    #[error("API key has been revoked")]
    KeyRevoked,
    #[error("API key has expired")]
    KeyExpired,
    #[error("Insufficient scope: missing {0}")]
    InsufficientScope(String),
    #[error("Key not found")]
    NotFound,
}

/// Key operation errors
#[derive(Debug, thiserror::Error)]
pub enum KeyError {
    #[error("Encryption failed")]
    EncryptionFailed,
    #[error("Decryption failed")]
    DecryptionFailed,
    #[error("Random generation failed")]
    RandomGenerationFailed,
    #[error("Invalid key format")]
    InvalidKeyFormat,
    #[error("Database error: {0}")]
    Database(String),
    #[error("Audit error: {0}")]
    Audit(String),
}

/// Key storage trait for database abstraction
#[async_trait::async_trait]
pub trait KeyStore: Send + Sync {
    /// Store a new key record
    async fn store_key(&self, record: &ApiKeyRecord) -> Result<(), KeyError>;
    /// Find keys by prefix
    async fn find_keys_by_prefix(&self, prefix: &str) -> Result<Vec<ApiKeyRecord>, KeyError>;
    /// Get key by ID
    async fn get_key(&self, key_id: Uuid) -> Result<Option<ApiKeyRecord>, KeyError>;
    /// Update key status
    async fn update_status(&self, key_id: Uuid, status: KeyStatus) -> Result<(), KeyError>;
    /// Update last used timestamp
    async fn update_last_used(&self, key_id: Uuid) -> Result<(), KeyError>;
    /// List all active keys
    async fn list_active_keys(&self) -> Result<Vec<ApiKeyRecord>, KeyError>;
    /// Set rotation info
    async fn set_rotation_info(&self, key_id: Uuid, info: KeyRotationInfo) -> Result<(), KeyError>;
}

/// Secure API Key Manager
pub struct ApiKeyManager {
    /// Key Encryption Key (loaded from HSM/KMS)
    kek: SecretVec<u8>,
    /// Database connection
    db: Arc<dyn KeyStore>,
    /// Audit logger
    audit: Arc<AuditLogger>,
    /// Key format prefix (live or test)
    key_prefix: String,
}

impl ApiKeyManager {
    /// Create a new API key manager
    ///
    /// # Arguments
    ///
    /// * `kek` - Key Encryption Key (should be loaded from HSM/KMS)
    /// * `db` - Key storage backend
    /// * `audit` - Audit logger
    /// * `is_production` - Whether to use live or test key prefix
    pub fn new(
        kek: SecretVec<u8>,
        db: Arc<dyn KeyStore>,
        audit: Arc<AuditLogger>,
        is_production: bool,
    ) -> Self {
        Self {
            kek,
            db,
            audit,
            key_prefix: if is_production { "rk_live_" } else { "rk_test_" }.to_string(),
        }
    }

    /// Generate a new API key
    ///
    /// Returns the plaintext key ONCE - it is never stored or recoverable.
    ///
    /// # Arguments
    ///
    /// * `tenant_id` - Organization/tenant identifier
    /// * `scopes` - Permission scopes for this key
    /// * `expires_in` - Optional expiration duration
    ///
    /// # Returns
    ///
    /// The generated key information including the plaintext key (only returned once)
    pub async fn generate_key(
        &self,
        tenant_id: Uuid,
        scopes: Vec<String>,
        expires_in: Option<Duration>,
    ) -> Result<GeneratedKey, KeyError> {
        // Generate cryptographically secure random key
        let mut key_bytes = generate_secure_random(32)?;
        let key_id = Uuid::new_v4();

        // Create the displayable key format: rk_live_<base58>
        let key_string = format!("{}{}", self.key_prefix, bs58::encode(&key_bytes).into_string());

        // Hash for validation (timing-safe comparison later)
        let key_hash = hash_key(&key_bytes);

        // Encrypt key material with KEK (envelope encryption)
        let (encrypted_key, nonce) = self.encrypt_key_material(&key_bytes)?;

        // Zeroize plaintext key bytes immediately
        zeroize_memory(&mut key_bytes);

        let expires_at = expires_in.map(|d| {
            Utc::now() + chrono::Duration::from_std(d).unwrap_or(chrono::Duration::days(365))
        });

        let record = ApiKeyRecord {
            key_id,
            tenant_id,
            key_prefix: key_string[..16.min(key_string.len())].to_string(),
            key_hash,
            encrypted_key,
            encryption_nonce: nonce,
            scopes: scopes.clone(),
            created_at: Utc::now(),
            expires_at,
            last_used_at: None,
            status: KeyStatus::Active,
            rotation: None,
        };

        // Store record (encrypted key only)
        self.db.store_key(&record).await?;

        // Audit log key creation
        self.audit
            .log(AuditEvent::new(
                AuditEventType::ApiKeyCreated,
                AuditActor::System,
            )
            .with_tenant(tenant_id)
            .with_key(key_id)
            .with_details(serde_json::json!({
                "scopes": scopes,
                "expires_at": expires_at,
                "key_prefix": record.key_prefix,
            }))
            .with_outcome(AuditOutcome::Success))
            .await
            .map_err(|e| KeyError::Audit(e.to_string()))?;

        Ok(GeneratedKey {
            key_id,
            api_key: SecretString::from(key_string),
            expires_at,
        })
    }

    /// Validate an API key
    ///
    /// Uses constant-time comparison to prevent timing attacks.
    ///
    /// # Arguments
    ///
    /// * `provided_key` - The API key string to validate
    /// * `required_scopes` - Scopes that must be present on the key
    /// * `client_ip` - Client IP address for audit logging
    ///
    /// # Returns
    ///
    /// Validated key information if successful
    pub async fn validate_key(
        &self,
        provided_key: &str,
        required_scopes: &[String],
        client_ip: std::net::IpAddr,
    ) -> Result<ValidatedKey, KeyValidationError> {
        // Extract key bytes from formatted string
        let key_bytes = self.parse_key_format(provided_key)?;

        // Hash the provided key
        let provided_hash = hash_key(&key_bytes);

        // Look up by prefix first (not sensitive)
        let prefix = &provided_key[..16.min(provided_key.len())];
        let candidates = self
            .db
            .find_keys_by_prefix(prefix)
            .await
            .map_err(|_| KeyValidationError::NotFound)?;

        // Find matching key using constant-time comparison
        let mut matched_record: Option<ApiKeyRecord> = None;
        for record in candidates {
            if constant_time_eq(&provided_hash, &record.key_hash) {
                matched_record = Some(record);
                break;
            }
        }

        let record = matched_record.ok_or(KeyValidationError::InvalidKey)?;

        // Check key status
        match record.status {
            KeyStatus::Revoked => {
                self.audit_failed_auth(&record, client_ip, "Key revoked")
                    .await;
                return Err(KeyValidationError::KeyRevoked);
            }
            KeyStatus::Expired => {
                self.audit_failed_auth(&record, client_ip, "Key expired")
                    .await;
                return Err(KeyValidationError::KeyExpired);
            }
            KeyStatus::Active | KeyStatus::Rotating => {}
        }

        // Check expiration
        if let Some(expires_at) = record.expires_at {
            if Utc::now() > expires_at {
                self.audit_failed_auth(&record, client_ip, "Key past expiration")
                    .await;
                return Err(KeyValidationError::KeyExpired);
            }
        }

        // Check scopes
        for required in required_scopes {
            if !record.scopes.contains(required) {
                self.audit_failed_auth(
                    &record,
                    client_ip,
                    &format!("Missing scope: {}", required),
                )
                .await;
                return Err(KeyValidationError::InsufficientScope(required.clone()));
            }
        }

        // Update last used timestamp
        let _ = self.db.update_last_used(record.key_id).await;

        // Audit successful auth
        let _ = self
            .audit
            .log(
                AuditEvent::new(AuditEventType::ApiKeyValidated, AuditActor::ApiKey(record.key_id))
                    .with_tenant(record.tenant_id)
                    .with_key(record.key_id)
                    .with_details(serde_json::json!({
                        "scopes_checked": required_scopes,
                    }))
                    .with_ip(client_ip)
                    .with_outcome(AuditOutcome::Success),
            )
            .await;

        Ok(ValidatedKey {
            key_id: record.key_id,
            tenant_id: record.tenant_id,
            scopes: record.scopes,
        })
    }

    /// Revoke an API key
    pub async fn revoke_key(
        &self,
        key_id: Uuid,
        reason: &str,
        actor: AuditActor,
    ) -> Result<(), KeyError> {
        let record = self
            .db
            .get_key(key_id)
            .await?
            .ok_or(KeyError::Database("Key not found".to_string()))?;

        self.db.update_status(key_id, KeyStatus::Revoked).await?;

        self.audit
            .log(
                AuditEvent::new(AuditEventType::ApiKeyRevoked, actor)
                    .with_tenant(record.tenant_id)
                    .with_key(key_id)
                    .with_details(serde_json::json!({
                        "reason": reason,
                        "key_prefix": record.key_prefix,
                    }))
                    .with_outcome(AuditOutcome::Success),
            )
            .await
            .map_err(|e| KeyError::Audit(e.to_string()))?;

        Ok(())
    }

    /// Get key by ID
    pub async fn get_key(&self, key_id: Uuid) -> Result<ApiKeyRecord, KeyError> {
        self.db
            .get_key(key_id)
            .await?
            .ok_or(KeyError::Database("Key not found".to_string()))
    }

    /// Update key status
    pub async fn update_status(&self, key_id: Uuid, status: KeyStatus) -> Result<(), KeyError> {
        self.db.update_status(key_id, status).await
    }

    /// Set rotation info
    pub async fn set_rotation_info(
        &self,
        key_id: Uuid,
        info: KeyRotationInfo,
    ) -> Result<(), KeyError> {
        self.db.set_rotation_info(key_id, info).await
    }

    /// List all active keys
    pub async fn list_active_keys(&self) -> Result<Vec<ApiKeyRecord>, KeyError> {
        self.db.list_active_keys().await
    }

    /// Parse key format and extract raw bytes
    fn parse_key_format(&self, key: &str) -> Result<Vec<u8>, KeyValidationError> {
        // Check prefix
        if !key.starts_with("rk_live_") && !key.starts_with("rk_test_") {
            return Err(KeyValidationError::InvalidKey);
        }

        // Extract base58 portion
        let encoded = &key[8..];
        bs58::decode(encoded)
            .into_vec()
            .map_err(|_| KeyValidationError::InvalidKey)
    }

    /// Encrypt key material using envelope encryption (AES-256-GCM)
    fn encrypt_key_material(&self, plaintext: &[u8]) -> Result<(Vec<u8>, [u8; 12]), KeyError> {
        use aes_gcm::{aead::Aead, Aes256Gcm, KeyInit, Nonce};

        let cipher = Aes256Gcm::new_from_slice(self.kek.expose_secret())
            .map_err(|_| KeyError::EncryptionFailed)?;

        let nonce_bytes = generate_secure_random(12)?;
        let nonce = Nonce::from_slice(&nonce_bytes);

        let ciphertext = cipher
            .encrypt(nonce, plaintext)
            .map_err(|_| KeyError::EncryptionFailed)?;

        let mut nonce_arr = [0u8; 12];
        nonce_arr.copy_from_slice(&nonce_bytes);

        Ok((ciphertext, nonce_arr))
    }

    /// Audit failed authentication attempt
    async fn audit_failed_auth(&self, record: &ApiKeyRecord, ip: std::net::IpAddr, reason: &str) {
        let _ = self
            .audit
            .log(
                AuditEvent::new(
                    AuditEventType::ApiKeyValidationFailed,
                    AuditActor::Anonymous,
                )
                .with_tenant(record.tenant_id)
                .with_key(record.key_id)
                .with_details(serde_json::json!({
                    "reason": reason,
                    "key_prefix": record.key_prefix,
                }))
                .with_ip(ip)
                .with_outcome(AuditOutcome::Failure {
                    error_code: "AUTH_FAILED".to_string(),
                    error_message: reason.to_string(),
                }),
            )
            .await;
    }
}

/// Constant-time byte comparison
///
/// CRITICAL: This function MUST take the same time regardless of where
/// the arrays differ. Any timing variance allows gradual key discovery.
#[inline(never)]
fn constant_time_eq(a: &[u8; 32], b: &[u8; 32]) -> bool {
    let mut result: u8 = 0;
    for (x, y) in a.iter().zip(b.iter()) {
        result |= x ^ y;
    }
    // Compiler barrier to prevent optimization
    std::sync::atomic::compiler_fence(std::sync::atomic::Ordering::SeqCst);
    result == 0
}

/// Generate cryptographically secure random bytes
fn generate_secure_random(len: usize) -> Result<Vec<u8>, KeyError> {
    use rand::RngCore;
    let mut bytes = vec![0u8; len];
    rand::rngs::OsRng
        .try_fill_bytes(&mut bytes)
        .map_err(|_| KeyError::RandomGenerationFailed)?;
    Ok(bytes)
}

/// Securely zeroize memory containing sensitive data
fn zeroize_memory(data: &mut Vec<u8>) {
    use zeroize::Zeroize;
    data.zeroize();
}

/// Hash key for storage (one-way, for comparison only)
fn hash_key(key: &[u8]) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(key);
    hasher.finalize().into()
}

// Serde helpers for binary data
mod hex_array {
    use serde::{Deserialize, Deserializer, Serialize, Serializer};

    pub fn serialize<S>(data: &[u8; 32], serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        hex::encode(data).serialize(serializer)
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<[u8; 32], D::Error>
    where
        D: Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        let bytes = hex::decode(&s).map_err(serde::de::Error::custom)?;
        let mut arr = [0u8; 32];
        arr.copy_from_slice(&bytes);
        Ok(arr)
    }
}

mod hex_nonce {
    use serde::{Deserialize, Deserializer, Serialize, Serializer};

    pub fn serialize<S>(data: &[u8; 12], serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        hex::encode(data).serialize(serializer)
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<[u8; 12], D::Error>
    where
        D: Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        let bytes = hex::decode(&s).map_err(serde::de::Error::custom)?;
        let mut arr = [0u8; 12];
        arr.copy_from_slice(&bytes);
        Ok(arr)
    }
}

mod base64_vec {
    use base64::{engine::general_purpose::STANDARD, Engine};
    use serde::{Deserialize, Deserializer, Serialize, Serializer};

    pub fn serialize<S>(data: &[u8], serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        STANDARD.encode(data).serialize(serializer)
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
    where
        D: Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        STANDARD.decode(&s).map_err(serde::de::Error::custom)
    }
}

/// In-memory key store for testing
pub struct InMemoryKeyStore {
    keys: RwLock<HashMap<Uuid, ApiKeyRecord>>,
}

impl InMemoryKeyStore {
    pub fn new() -> Self {
        Self {
            keys: RwLock::new(HashMap::new()),
        }
    }
}

impl Default for InMemoryKeyStore {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait::async_trait]
impl KeyStore for InMemoryKeyStore {
    async fn store_key(&self, record: &ApiKeyRecord) -> Result<(), KeyError> {
        let mut keys = self.keys.write().await;
        keys.insert(record.key_id, record.clone());
        Ok(())
    }

    async fn find_keys_by_prefix(&self, prefix: &str) -> Result<Vec<ApiKeyRecord>, KeyError> {
        let keys = self.keys.read().await;
        Ok(keys
            .values()
            .filter(|k| k.key_prefix.starts_with(prefix))
            .cloned()
            .collect())
    }

    async fn get_key(&self, key_id: Uuid) -> Result<Option<ApiKeyRecord>, KeyError> {
        let keys = self.keys.read().await;
        Ok(keys.get(&key_id).cloned())
    }

    async fn update_status(&self, key_id: Uuid, status: KeyStatus) -> Result<(), KeyError> {
        let mut keys = self.keys.write().await;
        if let Some(key) = keys.get_mut(&key_id) {
            key.status = status;
        }
        Ok(())
    }

    async fn update_last_used(&self, key_id: Uuid) -> Result<(), KeyError> {
        let mut keys = self.keys.write().await;
        if let Some(key) = keys.get_mut(&key_id) {
            key.last_used_at = Some(Utc::now());
        }
        Ok(())
    }

    async fn list_active_keys(&self) -> Result<Vec<ApiKeyRecord>, KeyError> {
        let keys = self.keys.read().await;
        Ok(keys
            .values()
            .filter(|k| k.status == KeyStatus::Active)
            .cloned()
            .collect())
    }

    async fn set_rotation_info(&self, key_id: Uuid, info: KeyRotationInfo) -> Result<(), KeyError> {
        let mut keys = self.keys.write().await;
        if let Some(key) = keys.get_mut(&key_id) {
            key.rotation = Some(info);
        }
        Ok(())
    }
}

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

    #[test]
    fn test_constant_time_eq_equal() {
        let a = [1u8; 32];
        let b = [1u8; 32];
        assert!(constant_time_eq(&a, &b));
    }

    #[test]
    fn test_constant_time_eq_unequal() {
        let a = [1u8; 32];
        let mut b = [1u8; 32];
        b[31] = 2;
        assert!(!constant_time_eq(&a, &b));
    }

    #[test]
    fn test_hash_key_deterministic() {
        let key = b"test_key_material_here";
        let hash1 = hash_key(key);
        let hash2 = hash_key(key);
        assert_eq!(hash1, hash2);
    }

    #[test]
    fn test_hash_key_different_inputs() {
        let key1 = b"test_key_1_material_";
        let key2 = b"test_key_2_material_";
        let hash1 = hash_key(key1);
        let hash2 = hash_key(key2);
        assert_ne!(hash1, hash2);
    }

    #[test]
    fn test_secure_random_length() {
        let bytes = generate_secure_random(32).unwrap();
        assert_eq!(bytes.len(), 32);
    }

    #[test]
    fn test_secure_random_uniqueness() {
        let bytes1 = generate_secure_random(32).unwrap();
        let bytes2 = generate_secure_random(32).unwrap();
        assert_ne!(bytes1, bytes2);
    }
}