atlas-chat 0.1.0__py3-none-any.whl

atlas/modules/mcp_tools/token_storage.py

@@ -0,0 +1,556 @@
+"""Secure per-user token storage for MCP server authentication.
+
+This module provides encrypted storage for authentication tokens associated
+with MCP servers on a per-user basis. Supports multiple token types:
+- API keys
+- JWT tokens
+- Bearer tokens
+- OAuth access tokens
+
+Each user's tokens are isolated and encrypted using Fernet (AES-128-CBC).
+
+Key format: "{user_email}:{server_name}"
+
+Security considerations:
+- Tokens are encrypted at rest using a key derived from environment variable
+- Each user's tokens are stored separately (isolation by key)
+- Token expiration is tracked and validated
+- No plaintext tokens are logged
+
+Updated: 2025-01-21
+"""
+
+import base64
+import json
+import logging
+import threading
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+from cryptography.fernet import Fernet, InvalidToken
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+
+logger = logging.getLogger(__name__)
+
+
+class AuthenticationRequiredException(Exception):
+    """Exception raised when a user needs to authenticate with an MCP server.
+
+    This exception carries information needed to initiate the OAuth flow
+    so the frontend can automatically redirect the user to authenticate.
+    """
+
+    def __init__(
+        self,
+        server_name: str,
+        auth_type: str,
+        message: str = "Authentication required",
+        oauth_start_url: Optional[str] = None,
+    ):
+        super().__init__(message)
+        self.server_name = server_name
+        self.auth_type = auth_type  # "oauth", "jwt", "bearer", or "api_key"
+        self.oauth_start_url = oauth_start_url  # URL to start OAuth flow (if oauth)
+        self.message = message
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert exception info to a dict for frontend consumption."""
+        return {
+            "server_name": self.server_name,
+            "auth_type": self.auth_type,
+            "message": self.message,
+            "oauth_start_url": self.oauth_start_url,
+        }
+
+
+def _make_token_key(user_email: str, server_name: str) -> str:
+    """Create a storage key from user email and server name.
+
+    Args:
+        user_email: User's email address
+        server_name: Name of the MCP server
+
+    Returns:
+        Combined key in format "user_email:server_name"
+    """
+    # Normalize to lowercase for consistent lookups
+    return f"{user_email.lower()}:{server_name}"
+
+
+def _parse_token_key(key: str) -> Tuple[str, str]:
+    """Parse a storage key into user email and server name.
+
+    Args:
+        key: Combined key in format "user_email:server_name"
+
+    Returns:
+        Tuple of (user_email, server_name)
+
+    Raises:
+        ValueError: If key format is invalid
+    """
+    parts = key.split(":", 1)
+    if len(parts) != 2:
+        raise ValueError(f"Invalid token key format: {key}")
+    return parts[0], parts[1]
+
+
+@dataclass
+class StoredToken:
+    """Represents a stored authentication token."""
+
+    token_type: str  # "api_key", "bearer", "jwt", "oauth_access", "oauth_refresh"
+    token_value: str  # The actual token (will be encrypted at rest)
+    user_email: str  # User who owns this token
+    server_name: str  # MCP server this token is associated with
+    created_at: float  # Unix timestamp when token was stored
+    expires_at: Optional[float] = None  # Unix timestamp when token expires (if known)
+    scopes: Optional[str] = None  # OAuth scopes (space-separated)
+    refresh_token: Optional[str] = None  # OAuth refresh token (if available)
+    metadata: Optional[Dict[str, Any]] = None  # Additional metadata
+
+    def is_expired(self, buffer_seconds: int = 60) -> bool:
+        """Check if token is expired or will expire within buffer period."""
+        if self.expires_at is None:
+            return False  # No expiration set, assume valid
+        return time.time() >= (self.expires_at - buffer_seconds)
+
+    def time_until_expiry(self) -> Optional[float]:
+        """Get seconds until token expires, or None if no expiration."""
+        if self.expires_at is None:
+            return None
+        return max(0, self.expires_at - time.time())
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "StoredToken":
+        """Create from dictionary."""
+        return cls(**data)
+
+
+class MCPTokenStorage:
+    """Secure encrypted storage for per-user MCP authentication tokens.
+
+    Tokens are stored in an encrypted JSON file on disk, keyed by the
+    combination of user email and server name. The encryption key is
+    derived from the MCP_TOKEN_ENCRYPTION_KEY environment variable using
+    PBKDF2. If no key is set, a random key is generated (tokens will not
+    persist across restarts in this case).
+
+    Storage location: {storage_dir}/mcp_tokens.enc
+    """
+
+    # Salt for key derivation (constant, not secret)
+    _SALT = b"atlas-mcp-token-storage-v1"
+
+    def __init__(
+        self,
+        storage_dir: Optional[Path] = None,
+        encryption_key: Optional[str] = None,
+    ):
+        """Initialize token storage.
+
+        Args:
+            storage_dir: Directory to store encrypted tokens. Defaults to
+                        MCP_TOKEN_STORAGE_DIR or config/secure
+            encryption_key: Base64-encoded encryption key or passphrase.
+                           Defaults to MCP_TOKEN_ENCRYPTION_KEY setting.
+        """
+        # Import here to avoid circular imports
+        from atlas.modules.config.config_manager import get_app_settings
+        app_settings = get_app_settings()
+
+        self._storage_dir = storage_dir or self._get_storage_dir(app_settings)
+        self._storage_dir.mkdir(parents=True, exist_ok=True)
+        self._storage_file = self._storage_dir / "mcp_tokens.enc"
+
+        # Get or generate encryption key (prefer passed arg, then settings)
+        key_source = encryption_key or app_settings.mcp_token_encryption_key
+        if key_source:
+            self._fernet = self._derive_fernet(key_source)
+            logger.info("Token storage initialized with configured encryption key")
+        else:
+            # Generate ephemeral key (tokens won't persist across restarts)
+            ephemeral_key = Fernet.generate_key()
+            self._fernet = Fernet(ephemeral_key)
+            logger.warning(
+                "No MCP_TOKEN_ENCRYPTION_KEY set. Using ephemeral key - "
+                "tokens will not persist across application restarts."
+            )
+
+        # Thread lock for concurrent access
+        self._lock = threading.Lock()
+
+        # In-memory cache of decrypted tokens
+        # Key format: "user_email:server_name"
+        self._tokens: Dict[str, StoredToken] = {}
+        self._load_tokens()
+
+    def _get_storage_dir(self, app_settings) -> Path:
+        """Get storage directory from settings or default locations.
+
+        Args:
+            app_settings: AppSettings instance with mcp_token_storage_dir
+
+        Returns:
+            Path to storage directory
+        """
+        # Check configured directory first
+        if app_settings.mcp_token_storage_dir:
+            configured_path = Path(app_settings.mcp_token_storage_dir)
+            try:
+                configured_path.mkdir(parents=True, exist_ok=True)
+                logger.info(f"Using token storage directory: {configured_path}")
+                return configured_path
+            except (PermissionError, OSError) as e:
+                logger.warning(f"Cannot use MCP_TOKEN_STORAGE_DIR={app_settings.mcp_token_storage_dir}: {e}")
+
+        # Try project root locations
+        candidates = [
+            Path(__file__).parent.parent.parent.parent / "config" / "secure",
+            Path(__file__).parent.parent.parent.parent / "runtime" / "tokens",
+            Path.home() / ".atlas-ui" / "tokens",
+        ]
+        for candidate in candidates:
+            try:
+                candidate.mkdir(parents=True, exist_ok=True)
+                # Test write access
+                test_file = candidate / ".write_test"
+                test_file.write_text("test")
+                test_file.unlink()
+                return candidate
+            except (PermissionError, OSError):
+                continue
+
+        # Fallback to temp directory
+        import tempfile
+        return Path(tempfile.gettempdir()) / "atlas-mcp-tokens"
+
+    def _derive_fernet(self, key_source: str) -> Fernet:
+        """Derive Fernet key from passphrase or base64 key."""
+        try:
+            # Try to use as direct Fernet key (base64-encoded 32 bytes)
+            return Fernet(key_source.encode())
+        except (ValueError, Exception):
+            # Derive key from passphrase using PBKDF2
+            kdf = PBKDF2HMAC(
+                algorithm=hashes.SHA256(),
+                length=32,
+                salt=self._SALT,
+                iterations=480000,  # OWASP recommended minimum
+            )
+            derived_key = base64.urlsafe_b64encode(
+                kdf.derive(key_source.encode())
+            )
+            return Fernet(derived_key)
+
+    def _load_tokens(self) -> None:
+        """Load and decrypt tokens from storage file."""
+        if not self._storage_file.exists():
+            self._tokens = {}
+            return
+
+        try:
+            encrypted_data = self._storage_file.read_bytes()
+            decrypted_data = self._fernet.decrypt(encrypted_data)
+            tokens_dict = json.loads(decrypted_data.decode())
+
+            self._tokens = {
+                key: StoredToken.from_dict(token_data)
+                for key, token_data in tokens_dict.items()
+            }
+            logger.info(f"Loaded {len(self._tokens)} encrypted tokens from storage")
+
+        except InvalidToken:
+            logger.error(
+                "Failed to decrypt token storage - encryption key may have changed. "
+                "Tokens will be reset."
+            )
+            self._tokens = {}
+        except (json.JSONDecodeError, KeyError, TypeError) as e:
+            logger.error(f"Corrupted token storage: {e}. Tokens will be reset.")
+            self._tokens = {}
+
+    def _save_tokens(self) -> None:
+        """Encrypt and save tokens to storage file."""
+        try:
+            tokens_dict = {
+                key: token.to_dict()
+                for key, token in self._tokens.items()
+            }
+            json_data = json.dumps(tokens_dict, indent=2)
+            encrypted_data = self._fernet.encrypt(json_data.encode())
+
+            # Atomic write: write to temp file then rename
+            temp_file = self._storage_file.with_suffix(".tmp")
+            temp_file.write_bytes(encrypted_data)
+            temp_file.rename(self._storage_file)
+
+            logger.debug(f"Saved {len(self._tokens)} encrypted tokens to storage")
+
+        except Exception as e:
+            logger.error(f"Failed to save tokens: {e}")
+            raise
+
+    def store_token(
+        self,
+        user_email: str,
+        server_name: str,
+        token_value: str,
+        token_type: str = "bearer",
+        expires_at: Optional[float] = None,
+        scopes: Optional[str] = None,
+        refresh_token: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> StoredToken:
+        """Store an authentication token for a user and MCP server.
+
+        Args:
+            user_email: User's email address
+            server_name: Name of the MCP server
+            token_value: The token value (JWT, access token, etc.)
+            token_type: Type of token ("bearer", "oauth_access", "jwt")
+            expires_at: Unix timestamp when token expires
+            scopes: OAuth scopes (space-separated string)
+            refresh_token: OAuth refresh token if available
+            metadata: Additional metadata to store
+
+        Returns:
+            The stored token object
+        """
+        token = StoredToken(
+            token_type=token_type,
+            token_value=token_value,
+            user_email=user_email.lower(),
+            server_name=server_name,
+            created_at=time.time(),
+            expires_at=expires_at,
+            scopes=scopes,
+            refresh_token=refresh_token,
+            metadata=metadata,
+        )
+
+        key = _make_token_key(user_email, server_name)
+        with self._lock:
+            self._tokens[key] = token
+            self._save_tokens()
+
+        from atlas.core.log_sanitizer import sanitize_for_logging
+        logger.info(
+            f"Stored {token_type} token for user and server '{sanitize_for_logging(server_name)}' "
+            f"(expires: {'never' if expires_at is None else time.ctime(expires_at)})"
+        )
+        return token
+
+    def get_token(self, user_email: str, server_name: str) -> Optional[StoredToken]:
+        """Get stored token for a user and MCP server.
+
+        Args:
+            user_email: User's email address
+            server_name: Name of the MCP server
+
+        Returns:
+            StoredToken if found, None otherwise
+        """
+        key = _make_token_key(user_email, server_name)
+        token = self._tokens.get(key)
+        if token is None:
+            return None
+
+        # Log warning if expired (but still return it - caller may want to refresh)
+        if token.is_expired():
+            logger.debug(f"Token for server '{server_name}' has expired")
+
+        return token
+
+    def get_valid_token(self, user_email: str, server_name: str) -> Optional[StoredToken]:
+        """Get stored token only if not expired.
+
+        Args:
+            user_email: User's email address
+            server_name: Name of the MCP server
+
+        Returns:
+            StoredToken if found and not expired, None otherwise
+        """
+        token = self.get_token(user_email, server_name)
+        if token is None or token.is_expired():
+            return None
+        return token
+
+    def remove_token(self, user_email: str, server_name: str) -> bool:
+        """Remove stored token for a user and MCP server.
+
+        Args:
+            user_email: User's email address
+            server_name: Name of the MCP server
+
+        Returns:
+            True if token was removed, False if not found
+        """
+        key = _make_token_key(user_email, server_name)
+        with self._lock:
+            if key in self._tokens:
+                del self._tokens[key]
+                self._save_tokens()
+                from atlas.core.log_sanitizer import sanitize_for_logging
+                logger.info(f"Removed token for server '{sanitize_for_logging(server_name)}'")
+                return True
+            return False
+
+    def get_user_tokens(self, user_email: str) -> Dict[str, StoredToken]:
+        """Get all tokens for a specific user.
+
+        Args:
+            user_email: User's email address
+
+        Returns:
+            Dictionary mapping server names to tokens
+        """
+        user_email_lower = user_email.lower()
+        return {
+            _parse_token_key(key)[1]: token  # Extract server_name from key
+            for key, token in self._tokens.items()
+            if token.user_email == user_email_lower
+        }
+
+    def get_user_auth_status(self, user_email: str) -> Dict[str, Dict[str, Any]]:
+        """Get authentication status for all servers for a user.
+
+        Returns metadata about tokens without revealing token values.
+
+        Args:
+            user_email: User's email address
+
+        Returns:
+            Dictionary mapping server names to auth status info
+        """
+        user_tokens = self.get_user_tokens(user_email)
+        return {
+            server_name: {
+                "authenticated": True,
+                "token_type": token.token_type,
+                "created_at": token.created_at,
+                "expires_at": token.expires_at,
+                "is_expired": token.is_expired(),
+                "time_until_expiry": token.time_until_expiry(),
+                "has_refresh_token": token.refresh_token is not None,
+                "scopes": token.scopes,
+            }
+            for server_name, token in user_tokens.items()
+        }
+
+    def list_all_tokens_metadata(self) -> List[Dict[str, Any]]:
+        """List metadata for all stored tokens (admin use).
+
+        Returns token metadata without revealing token values.
+
+        Returns:
+            List of token metadata dictionaries
+        """
+        return [
+            {
+                "user_email": token.user_email,
+                "server_name": token.server_name,
+                "token_type": token.token_type,
+                "created_at": token.created_at,
+                "expires_at": token.expires_at,
+                "is_expired": token.is_expired(),
+                "has_refresh_token": token.refresh_token is not None,
+                "scopes": token.scopes,
+            }
+            for token in self._tokens.values()
+        ]
+
+    def update_oauth_tokens(
+        self,
+        user_email: str,
+        server_name: str,
+        access_token: str,
+        expires_at: Optional[float] = None,
+        refresh_token: Optional[str] = None,
+        scopes: Optional[str] = None,
+    ) -> StoredToken:
+        """Update OAuth tokens after a refresh or new authorization.
+
+        Preserves existing metadata and refresh token if new one not provided.
+
+        Args:
+            user_email: User's email address
+            server_name: Name of the MCP server
+            access_token: New access token
+            expires_at: Unix timestamp when token expires
+            refresh_token: New refresh token (or None to keep existing)
+            scopes: OAuth scopes
+
+        Returns:
+            Updated StoredToken
+        """
+        existing = self.get_token(user_email, server_name)
+
+        return self.store_token(
+            user_email=user_email,
+            server_name=server_name,
+            token_value=access_token,
+            token_type="oauth_access",
+            expires_at=expires_at,
+            scopes=scopes or (existing.scopes if existing else None),
+            refresh_token=refresh_token or (existing.refresh_token if existing else None),
+            metadata=existing.metadata if existing else None,
+        )
+
+    def clear_user_tokens(self, user_email: str) -> int:
+        """Remove all tokens for a specific user.
+
+        Args:
+            user_email: User's email address
+
+        Returns:
+            Number of tokens removed
+        """
+        user_email_lower = user_email.lower()
+        with self._lock:
+            keys_to_remove = [
+                key for key, token in self._tokens.items()
+                if token.user_email == user_email_lower
+            ]
+
+            for key in keys_to_remove:
+                del self._tokens[key]
+
+            if keys_to_remove:
+                self._save_tokens()
+                logger.info(f"Cleared {len(keys_to_remove)} tokens for user")
+
+        return len(keys_to_remove)
+
+    def clear_all(self) -> int:
+        """Remove all stored tokens (admin use).
+
+        Returns:
+            Number of tokens removed
+        """
+        with self._lock:
+            count = len(self._tokens)
+            self._tokens.clear()
+            self._save_tokens()
+            logger.info(f"Cleared all {count} stored tokens")
+        return count
+
+
+# Global token storage instance (lazy initialization)
+_token_storage: Optional[MCPTokenStorage] = None
+
+
+def get_token_storage() -> MCPTokenStorage:
+    """Get the global token storage instance."""
+    global _token_storage
+    if _token_storage is None:
+        _token_storage = MCPTokenStorage()
+    return _token_storage

atlas/modules/prompts/prompt_provider.py

@@ -0,0 +1,130 @@
+"""Prompt provider module for loading and caching prompt templates.
+
+Centralizes prompt path resolution & template retrieval so core services stay
+focused on orchestration/business logic.
+"""
+from __future__ import annotations
+
+import logging
+import os
+from typing import Dict, Optional
+
+from atlas.modules.config import ConfigManager
+
+logger = logging.getLogger(__name__)
+
+
+class PromptProvider:
+    """Loads and caches prompt templates based on application configuration."""
+
+    def __init__(self, config_manager: ConfigManager):
+        self.config_manager = config_manager
+        self._cache: Dict[str, str] = {}
+        # Resolve base path (relative paths resolved against repo root)
+        app_settings = self.config_manager.app_settings
+        repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
+        base_candidate = app_settings.prompt_base_path
+        if not os.path.isabs(base_candidate):
+            self.base_path = os.path.join(repo_root, base_candidate)
+        else:
+            self.base_path = base_candidate
+
+    def _load_template(self, filename: str) -> Optional[str]:
+        cache_key = filename
+        if cache_key in self._cache:
+            return self._cache[cache_key]
+        path = os.path.join(self.base_path, filename)
+        if not os.path.exists(path):
+            logger.warning("Prompt template not found: %s", path)
+            return None
+        try:
+            with open(path, "r", encoding="utf-8") as f:
+                content = f.read()
+            self._cache[cache_key] = content
+            return content
+        except Exception as e:  # pragma: no cover
+            logger.error("Failed reading prompt template %s: %s", path, e)
+            return None
+
+    def get_tool_synthesis_prompt(self, user_question: str) -> Optional[str]:
+        """Return formatted tool synthesis prompt or None if unavailable."""
+        filename = self.config_manager.app_settings.tool_synthesis_prompt_filename
+        template = self._load_template(filename)
+        if not template:
+            return None
+        try:
+            return template.format(user_question=user_question.strip())
+        except Exception as e:  # pragma: no cover - safeguard
+            logger.warning("Formatting tool synthesis prompt failed: %s", e)
+            return None
+
+    def get_agent_reason_prompt(
+        self,
+        user_question: str,
+        files_manifest: Optional[str] = None,
+        last_observation: Optional[str] = None,
+    ) -> Optional[str]:
+        """Return formatted agent reason prompt text or None if unavailable.
+
+        Expects template placeholders: {user_question}, {files_manifest}, {last_observation}
+        Missing values are rendered as empty strings.
+        """
+        filename = self.config_manager.app_settings.agent_reason_prompt_filename
+        template = self._load_template(filename)
+        if not template:
+            return None
+        try:
+            return template.format(
+                user_question=(user_question or "").strip(),
+                files_manifest=(files_manifest or ""),
+                last_observation=(last_observation or ""),
+            )
+        except Exception as e:  # pragma: no cover
+            logger.warning("Formatting agent reason prompt failed: %s", e)
+            return None
+
+    def get_agent_observe_prompt(
+        self,
+        user_question: str,
+        tool_summaries: str,
+        step: int,
+    ) -> Optional[str]:
+        """Return formatted agent observe prompt text or None if unavailable.
+
+        Expects template placeholders: {user_question}, {tool_summaries}, {step}
+        """
+        filename = self.config_manager.app_settings.agent_observe_prompt_filename
+        template = self._load_template(filename)
+        if not template:
+            return None
+        try:
+            return template.format(
+                user_question=(user_question or "").strip(),
+                tool_summaries=(tool_summaries or ""),
+                step=step,
+            )
+        except Exception as e:  # pragma: no cover
+            logger.warning("Formatting agent observe prompt failed: %s", e)
+            return None
+
+    def get_system_prompt(self, user_email: Optional[str] = None) -> Optional[str]:
+        """Return formatted system prompt text or None if unavailable.
+
+        Expects template placeholder: {user_email}
+        Missing values are rendered as empty strings.
+        """
+        filename = self.config_manager.app_settings.system_prompt_filename
+        template = self._load_template(filename)
+        if not template:
+            return None
+        try:
+            return template.format(
+                user_email=(user_email or ""),
+            )
+        except Exception as e:  # pragma: no cover
+            logger.warning("Formatting system prompt failed: %s", e)
+            return None
+
+    def clear_cache(self) -> None:
+        """Clear in-memory prompt cache (e.g., after config reload)."""
+        self._cache.clear()