security/secret_manager.py

"""
Secret Manager for AegisLM

Provides secure secret management and retrieval.
Loads secrets from environment variables at runtime.
"""

import os
from typing import Any, Dict, Optional
from functools import lru_cache

from pydantic import BaseModel


class SecretManagerConfig(BaseModel):
    """Configuration for secret manager."""
    # JWT secrets
    jwt_secret_key: str
    jwt_algorithm: str = "HS256"
    jwt_expiration_hours: int = 24
    
    # Database
    database_url: str
    
    # Encryption
    encryption_key: str  # AES-256 key (base64 encoded)
    
    # Worker authentication
    worker_secret_key: str
    
    # API keys for external services (optional)
    hf_token: Optional[str] = None
    openai_api_key: Optional[str] = None
    anthropic_api_key: Optional[str] = None
    
    # Additional secrets
    secret_prefix: str = "AEGISLM_"


class SecretManager:
    """
    Centralized secret management.
    
    Loads secrets from environment variables at runtime.
    Secrets should NEVER be hardcoded in source files.
    """
    
    _instance: Optional["SecretManager"] = None
    _config: Optional[SecretManagerConfig] = None
    
    def __new__(cls) -> "SecretManager":
        """Singleton pattern for secret manager."""
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance
    
    def __init__(self):
        """Initialize the secret manager."""
        if self._config is None:
            self._load_config()
    
    def _load_config(self) -> None:
        """Load configuration from environment variables."""
        prefix = os.getenv("AEGISLM_SECRET_PREFIX", "AEGISLM_")
        
        # Check production mode first - this enables strict enforcement
        self._is_production = os.getenv("AEGISLM_ENV", "development").lower() == "production"
        
        # Required secrets - these MUST be set in environment
        jwt_secret = os.getenv(f"{prefix}JWT_SECRET_KEY")
        if not jwt_secret:
            jwt_secret = os.getenv("JWT_SECRET_KEY")
            if not jwt_secret:
                if self._is_production:
                    raise ValueError(
                        "CRITICAL: JWT_SECRET_KEY must be set in environment variables. "
                        "Production mode requires AEGISLM_JWT_SECRET_KEY to be set. "
                        "Cannot start without secure JWT secret."
                    )
                else:
                    jwt_secret = os.getenv("JWT_SECRET_KEY", "dev-secret-key-change-in-production")
        
        # Check for default dev key in production mode - FAIL FAST
        if self._is_production and jwt_secret == "dev-secret-key-change-in-production":
            raise ValueError(
                "CRITICAL: Production mode detected but using default dev secret key. "
                "This is a security violation. Set AEGISLM_JWT_SECRET_KEY to a secure value. "
                "Refusing to start in production with insecure default key."
            )
        
        db_url = os.getenv(f"{prefix}DATABASE_URL")
        if not db_url:
            db_url = os.getenv("DATABASE_URL")
            if not db_url:
                if self._is_production:
                    raise ValueError(
                        "CRITICAL: DATABASE_URL must be set in environment variables. "
                        "Production mode requires PostgreSQL database. Cannot start without DATABASE_URL."
                    )
                else:
                    db_url = "sqlite+aiosqlite:///./aegislm_dev.db"
        
        # Check for SQLite in production - FAIL FAST
        if self._is_production and "sqlite" in db_url.lower():
            raise ValueError(
                "CRITICAL: Production mode does not support SQLite. "
                "Use PostgreSQL with asyncpg driver: postgresql+asyncpg://user:pass@host:5432/db"
            )
        
        # Check Redis URL for production
        redis_url = os.getenv(f"{prefix}REDIS_URL")
        if not redis_url:
            redis_url = os.getenv("REDIS_URL")
            if not redis_url:
                if self._is_production:
                    raise ValueError(
                        "CRITICAL: REDIS_URL must be set in environment variables. "
                        "Production mode requires Redis for rate limiting. Cannot start without REDIS_URL."
                    )
                else:
                    redis_url = "redis://localhost:6379/0"
        
        if self._is_production and redis_url == "redis://localhost:6379/0":
            # This is just a warning - Redis might be available
            pass
        
        encryption_key = os.getenv(f"{prefix}ENCRYPTION_KEY")
        if not encryption_key:
            import warnings
            warnings.warn(
                "ENCRYPTION_KEY not set. Using default (insecure). "
                "Set AEGISLM_ENCRYPTION_KEY for production.",
                DeprecationWarning
            )
            encryption_key = "default-encryption-key-32-bytes!!"  # 32 bytes for AES-256
        
        worker_secret = os.getenv(f"{prefix}WORKER_SECRET_KEY")
        if not worker_secret:
            worker_secret = os.getenv("WORKER_SECRET_KEY", "worker-secret-key-change-in-production")
        
        self._config = SecretManagerConfig(
            jwt_secret_key=jwt_secret,
            database_url=db_url,
            encryption_key=encryption_key,
            worker_secret_key=worker_secret,
            hf_token=os.getenv(f"{prefix}HF_TOKEN") or os.getenv("HF_TOKEN"),
            openai_api_key=os.getenv(f"{prefix}OPENAI_API_KEY") or os.getenv("OPENAI_API_KEY"),
            anthropic_api_key=os.getenv(f"{prefix}ANTHROPIC_API_KEY") or os.getenv("ANTHROPIC_API_KEY"),
            secret_prefix=prefix,
        )
    
    @property
    def config(self) -> SecretManagerConfig:
        """Get the secret manager configuration."""
        if self._config is None:
            self._load_config()
        return self._config
    
    def get_jwt_secret(self) -> str:
        """Get JWT secret key."""
        return self.config.jwt_secret_key
    
    def get_jwt_algorithm(self) -> str:
        """Get JWT algorithm."""
        return self.config.jwt_algorithm
    
    def get_jwt_expiration_hours(self) -> int:
        """Get JWT token expiration in hours."""
        return self.config.jwt_expiration_hours
    
    def get_database_url(self) -> str:
        """Get database URL."""
        return self.config.database_url
    
    def get_encryption_key(self) -> str:
        """Get encryption key."""
        return self.config.encryption_key
    
    def get_worker_secret(self) -> str:
        """Get worker secret key."""
        return self.config.worker_secret_key
    
    def get_hf_token(self) -> Optional[str]:
        """Get HuggingFace token."""
        return self.config.hf_token
    
    def get_openai_api_key(self) -> Optional[str]:
        """Get OpenAI API key."""
        return self.config.openai_api_key
    
    def get_anthropic_api_key(self) -> Optional[str]:
        """Get Anthropic API key."""
        return self.config.anthropic_api_key
    
    def get_secret(self, name: str, default: Optional[str] = None) -> Optional[str]:
        """
        Get a secret by name.
        
        Args:
            name: Secret name (without prefix)
            default: Default value if not found
            
        Returns:
            Secret value or default
        """
        prefix = self.config.secret_prefix
        full_name = f"{prefix}{name}"
        
        # Try with prefix
        value = os.getenv(full_name)
        if value:
            return value
        
        # Try without prefix (for backwards compatibility)
        value = os.getenv(name)
        if value:
            return value
        
        return default
    
    def get_required_secret(self, name: str) -> str:
        """
        Get a required secret by name.
        
        Args:
            name: Secret name (without prefix)
            
        Returns:
            Secret value
            
        Raises:
            ValueError: If secret is not found
        """
        value = self.get_secret(name)
        if not value:
            raise ValueError(f"Required secret '{name}' not found in environment variables")
        return value
    
    def is_production_mode(self) -> bool:
        """Check if running in production mode."""
        return os.getenv("AEGISLM_ENV", "development").lower() == "production"
    
    def validate_secrets(self) -> Dict[str, bool]:
        """
        Validate that all required secrets are set.
        
        Returns:
            Dictionary of secret names and whether they're set
        """
        required_secrets = [
            "JWT_SECRET_KEY",
            "DATABASE_URL",
            "ENCRYPTION_KEY",
            "WORKER_SECRET_KEY",
        ]
        
        return {
            secret: self.get_secret(secret) is not None
            for secret in required_secrets
        }


@lru_cache()
def get_secret_manager() -> SecretManager:
    """
    Get the singleton secret manager instance.
    
    Returns:
        SecretManager instance
    """
    return SecretManager()


# Convenience functions
def get_jwt_secret() -> str:
    """Get JWT secret key."""
    return get_secret_manager().get_jwt_secret()


def get_jwt_algorithm() -> str:
    """Get JWT algorithm."""
    return get_secret_manager().get_jwt_algorithm()


def get_database_url() -> str:
    """Get database URL."""
    return get_secret_manager().get_database_url()


def get_encryption_key() -> str:
    """Get encryption key."""
    return get_secret_manager().get_encryption_key()


def get_worker_secret() -> str:
    """Get worker secret key."""
    return get_secret_manager().get_worker_secret()