atlas-chat 0.1.0__py3-none-any.whl

atlas/core/auth.py

@@ -0,0 +1,205 @@
+"""Authentication and authorization module."""
+
+import hmac
+import logging
+import re
+from datetime import datetime, timedelta
+from typing import Dict, Optional, Tuple
+
+import httpx
+import jwt
+
+from atlas.modules.config.config_manager import config_manager
+
+logger = logging.getLogger(__name__)
+
+# Cache with TTL for ALB public keys: {(kid, region): (key, expiry_time)}
+_alb_key_cache: Dict[Tuple[str, str], Tuple[str, datetime]] = {}
+
+
+async def is_user_in_group(user_id: str, group_id: str) -> bool:
+    """
+    Check if a user is in a specified group.
+
+    This function first checks for a configured external authorization endpoint.
+    If available, it makes an HTTP request to check group membership.
+    If not configured, it falls back to a mock implementation for local development.
+
+    Args:
+        user_id: User email/identifier.
+        group_id: Group identifier.
+
+    Returns:
+        True if the user is in the group, False otherwise.
+    """
+    app_settings = config_manager.app_settings
+    auth_url = app_settings.auth_group_check_url
+    api_key = app_settings.auth_group_check_api_key
+
+    if auth_url and api_key:
+        # Use the external HTTP endpoint for authorization
+        try:
+            async with httpx.AsyncClient() as client:
+                headers = {"Authorization": f"Bearer {api_key}"}
+                payload = {"user_id": user_id, "group_id": group_id}
+                response = await client.post(auth_url, json=payload, headers=headers, timeout=5.0)
+                response.raise_for_status()
+                # Assuming the endpoint returns a simple JSON like {"is_member": true}
+                return response.json().get("is_member", False)
+        except httpx.RequestError as e:
+            logger.error(f"HTTP request to auth endpoint failed: {e}", exc_info=True)
+            return False
+        except Exception as e:
+            logger.error(f"Error during external auth check: {e}", exc_info=True)
+            return False
+    else:
+        # Everybody is in the users group by default
+        if (group_id == "users"):
+            return True
+        # Fallback to mock implementation if no external endpoint is configured
+        if (app_settings.debug_mode and
+                user_id == app_settings.test_user and
+                group_id == app_settings.admin_group):
+            return True
+
+        mock_groups = {
+            "test@test.com": ["users", "mcp_basic", "admin"],
+            "user@example.com": ["users", "mcp_basic"],
+            "admin@example.com": ["admin", "users", "mcp_basic", "mcp_advanced"]
+        }
+        user_groups = mock_groups.get(user_id, [])
+        return group_id in user_groups
+
+
+def _get_alb_public_key(kid: str, aws_region: str) -> Optional[str]:
+    """
+    Fetch and cache AWS ALB public key by key ID.
+
+    Caching reduces latency and API calls since AWS ALB rotates keys infrequently.
+    Cache has a 1-hour TTL to handle key rotation.
+
+    Args:
+        kid: Key ID from JWT header
+        aws_region: AWS region (e.g., 'us-east-1')
+
+    Returns:
+        Public key string, or None if fetch fails
+    """
+    # Security: Validate inputs to prevent URL injection and cache poisoning attacks
+    # kid and region are used in URL construction, so strict validation is critical
+    if not re.match(r'^[a-zA-Z0-9\-]+$', kid):
+        logger.error(f"Invalid kid format: {kid}")
+        return None
+    if not re.match(r'^[a-z]{2}-[a-z]+-\d+$', aws_region):
+        logger.error(f"Invalid AWS region format: {aws_region}")
+        return None
+
+    # Security: TTL-based cache (1 hour) allows key rotation and prevents stale keys
+    # if AWS rotates keys or a key is compromised
+    cache_key = (kid, aws_region)
+    now = datetime.utcnow()
+    if cache_key in _alb_key_cache:
+        cached_key, expiry = _alb_key_cache[cache_key]
+        if now < expiry:
+            return cached_key
+        else:
+            # Expired, remove from cache
+            del _alb_key_cache[cache_key]
+
+    url = f'https://public-keys.auth.elb.{aws_region}.amazonaws.com/{kid}'
+    try:
+        response = httpx.get(url, timeout=5.0)
+        response.raise_for_status()
+        pub_key = response.text
+
+        # Cache with 1-hour TTL
+        expiry = now + timedelta(hours=1)
+        _alb_key_cache[cache_key] = (pub_key, expiry)
+
+        return pub_key
+    except httpx.HTTPStatusError as e:
+        logger.error(f"HTTP error fetching ALB public key from {url}: {e.response.status_code}")
+        return None
+    except httpx.RequestError as e:
+        logger.error(f"Error fetching ALB public key from {url}: {e}")
+        return None
+
+
+def get_user_from_aws_alb_jwt(encoded_jwt, expected_alb_arn, aws_region):
+    """
+    Validates the AWS ALB JWT and parses the email address from the payload.
+
+    Args:
+        encoded_jwt (str): The JWT from the x-amzn-oidc-data header.
+        expected_alb_arn (str): The ARN of your Application Load Balancer.
+        aws_region (str): The AWS region where your ALB is located (e.g., 'us-east-1').
+
+    Returns:
+        str: The user's email address, or None if validation fails.
+    """
+    if not encoded_jwt:
+        return None
+    try:
+        # Step 1: Decode the JWT header to get the key ID (kid) and signer using PyJWT
+        header = jwt.get_unverified_header(encoded_jwt)
+        kid = header.get('kid')
+        received_alb_arn = header.get('signer')
+
+        if not kid:
+            logger.error("Error: 'kid' not found in JWT header")
+            return None
+
+        # Step 2: Validate the signer matches the expected ALB ARN
+        # Security: hmac.compare_digest prevents timing attacks that could reveal the ARN
+        if not received_alb_arn or not hmac.compare_digest(received_alb_arn, expected_alb_arn):
+            logger.error(f"Error: Invalid signer ARN. Expected {expected_alb_arn}, got {received_alb_arn}")
+            return None
+
+        # Step 3: Get the public key from the regional endpoint (with caching)
+        pub_key = _get_alb_public_key(kid, aws_region)
+        if not pub_key:
+            logger.error("Error: Failed to fetch ALB public key")
+            return None
+
+        # Step 4: Validate the signature and claims using PyJWT
+        # The decode method handles signature verification and standard claims (like expiration)
+        # The ALB uses ES256 algorithm
+        payload = jwt.decode(
+            encoded_jwt,
+            pub_key,
+            algorithms=['ES256'],
+            # Optional: Add audience or issuer validation if needed, though ALB handles most standard claims validation
+            options={"verify_aud": False, "verify_iss": False}
+        )
+
+        # Step 5: Extract the email address from the payload
+        email_address = payload.get('email')
+        if email_address:
+            # Security: Validate email format to prevent injection attacks and ensure
+            # the email claim contains a properly formatted email address
+            email_pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
+            if not isinstance(email_address, str) or not re.match(email_pattern, email_address):
+                logger.error(f"Error: Invalid email format in JWT payload: {email_address}")
+                return None
+            logger.debug("Successfully authenticated user via AWS ALB JWT")
+            return email_address
+        else:
+            logger.error("Error: 'email' claim not found in JWT payload")
+            return None
+
+    except jwt.ExpiredSignatureError:
+        logger.error("Error: Token has expired")
+        return None
+    except jwt.InvalidTokenError as e:
+        logger.error(f"Error: Invalid token - {e}")
+        return None
+    except Exception as e:
+        logger.error(f"An unexpected error occurred: {e}")
+        return None
+
+
+def get_user_from_header(x_email_header: Optional[str]) -> Optional[str]:
+    """Extract user email from authentication header value."""
+    if not x_email_header:
+        return None
+    return x_email_header.strip()

atlas/core/authorization_manager.py

@@ -0,0 +1,27 @@
+"""Authorization utilities for managing access to resources."""
+
+import logging
+from typing import Awaitable, Callable
+
+from atlas.modules.config.config_manager import get_app_settings
+
+logger = logging.getLogger(__name__)
+
+AuthCheckFunc = Callable[[str, str], Awaitable[bool]]
+
+
+class AuthorizationManager:
+    """Manages authorization logic for admin access."""
+
+    def __init__(self, auth_check_func: AuthCheckFunc):
+        self.auth_check_func = auth_check_func
+        self.app_settings = get_app_settings()
+
+    async def is_admin(self, user_email: str) -> bool:
+        """Check if a user has admin privileges."""
+        return await self.auth_check_func(user_email, self.app_settings.admin_group)
+
+
+def create_authorization_manager(auth_check_func: AuthCheckFunc) -> AuthorizationManager:
+    """Factory function to create an AuthorizationManager."""
+    return AuthorizationManager(auth_check_func)

atlas/core/capabilities.py

@@ -0,0 +1,123 @@
+"""
+Capability-token utilities for secure, headless access to resources.
+
+Provides short-lived HMAC-signed tokens suitable for embedding in URLs,
+primarily for file downloads by tools that don't carry session cookies.
+"""
+
+import base64
+import hmac
+import json
+import logging
+import time
+from hashlib import sha256
+from typing import Any, Dict, Optional
+
+from atlas.modules.config import config_manager
+
+logger = logging.getLogger(__name__)
+
+
+def _b64url_encode(data: bytes) -> str:
+    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
+
+
+def _b64url_decode(data: str) -> bytes:
+    padding = "=" * (-len(data) % 4)
+    return base64.urlsafe_b64decode((data + padding).encode("ascii"))
+
+
+def _get_secret() -> bytes:
+    """Get the capability token secret as bytes.
+
+    Order of precedence:
+    - App settings (config manager)
+    - Fallback development secret (unsafe for production)
+    """
+    try:
+        settings = config_manager.app_settings
+        if getattr(settings, "capability_token_secret", None):
+            return settings.capability_token_secret.encode("utf-8")
+    except Exception:
+        # Config not ready; continue to fallback with a dev secret.
+        logger.debug("Capability token secret not available; using fallback dev secret.")
+
+    logger.warning("Using fallback dev capability token secret. Set CAPABILITY_TOKEN_SECRET for security.")
+    return b"dev-capability-secret"
+
+
+def _get_default_ttl_seconds() -> int:
+    try:
+        settings = config_manager.app_settings
+        ttl = getattr(settings, "capability_token_ttl_seconds", None)
+        if isinstance(ttl, int) and ttl > 0:
+            return ttl
+    except Exception:
+        logger.debug("Capability token TTL not available; using default TTL.")
+    return 3600
+
+
+def generate_file_token(user_email: str, file_key: str, ttl_seconds: Optional[int] = None) -> str:
+    """Generate a short-lived token authorizing access to a file key for a user."""
+    exp = int(time.time()) + (ttl_seconds or _get_default_ttl_seconds())
+    payload = {"u": user_email, "k": file_key, "e": exp}
+    payload_bytes = json.dumps(payload, separators=(",", ":")).encode("utf-8")
+    body = _b64url_encode(payload_bytes)
+    sig = hmac.new(_get_secret(), body.encode("ascii"), sha256).digest()
+    return f"{body}.{_b64url_encode(sig)}"
+
+
+def verify_file_token(token: str) -> Optional[Dict[str, Any]]:
+    """Verify a file token and return claims if valid, else None."""
+    try:
+        body, sig_b64 = token.split(".", 1)
+        expected_sig = hmac.new(_get_secret(), body.encode("ascii"), sha256).digest()
+        given_sig = _b64url_decode(sig_b64)
+        if not hmac.compare_digest(expected_sig, given_sig):
+            return None
+
+        claims = json.loads(_b64url_decode(body).decode("utf-8"))
+        if int(claims.get("e", 0)) < int(time.time()):
+            return None
+        # Ensure required claims exist
+        if not claims.get("u") or not claims.get("k"):
+            return None
+        return claims
+    except Exception:
+        return None
+
+
+def create_download_url(file_key: str, user_email: Optional[str]) -> str:
+    """Create a download URL for a given file key, optionally with a token.
+
+    If BACKEND_PUBLIC_URL is configured, returns an absolute URL that remote MCP servers
+    can access. Otherwise, returns a relative URL (only works for local/stdio servers).
+
+    Args:
+        file_key: S3 key of the file to download
+        user_email: User email for token generation
+
+    Returns:
+        Download URL (absolute if BACKEND_PUBLIC_URL configured, relative otherwise)
+    """
+    # Build relative path with token
+    if user_email:
+        token = generate_file_token(user_email, file_key)
+        relative_path = f"/api/files/download/{file_key}?token={token}"
+    else:
+        # Fallback: no user context available
+        relative_path = f"/api/files/download/{file_key}"
+
+    # Check if we should use absolute URLs for remote MCP server access
+    try:
+        settings = config_manager.app_settings
+        backend_public_url = getattr(settings, "backend_public_url", None)
+        if backend_public_url:
+            # Strip trailing slash from base URL and combine with relative path
+            base = backend_public_url.rstrip("/")
+            return f"{base}{relative_path}"
+    except Exception as e:
+        logger.debug(f"Could not check backend_public_url config: {e}")
+
+    # Return relative URL as default
+    return relative_path

atlas/core/compliance.py

@@ -0,0 +1,215 @@
+"""
+Compliance level management and validation.
+
+Loads compliance level definitions from compliance-levels.json and provides
+validation and allowlist checking.
+"""
+
+import json
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Set
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ComplianceLevel:
+    """Represents a single compliance level definition."""
+    name: str
+    description: str
+    aliases: List[str]
+    allowed_with: List[str]  # List of compliance levels that can be used together
+
+
+class ComplianceLevelManager:
+    """Manages compliance level definitions and validation."""
+
+    def __init__(self, config_path: Optional[Path] = None):
+        """Initialize the compliance level manager.
+
+        Args:
+            config_path: Path to compliance-levels.json. If None, uses default location.
+        """
+        self.levels: Dict[str, ComplianceLevel] = {}
+        self.mode: str = "explicit_allowlist"
+        self._name_to_canonical: Dict[str, str] = {}  # Maps aliases to canonical names
+
+        if config_path is None:
+            # Try to find config in standard locations
+            backend_root = Path(__file__).parent.parent
+            project_root = backend_root.parent
+
+            search_paths = [
+                project_root / "config" / "overrides" / "compliance-levels.json",
+                project_root / "config" / "defaults" / "compliance-levels.json",
+                backend_root / "configfilesadmin" / "compliance-levels.json",
+                backend_root / "configfiles" / "compliance-levels.json",
+            ]
+
+            for path in search_paths:
+                if path.exists():
+                    config_path = path
+                    break
+
+        if config_path and config_path.exists():
+            self._load_config(config_path)
+        else:
+            logger.warning("No compliance-levels.json found, using permissive validation")
+
+    def _load_config(self, config_path: Path):
+        """Load compliance level configuration from JSON file."""
+        try:
+            with open(config_path, 'r', encoding='utf-8') as f:
+                config = json.load(f)
+
+            self.mode = config.get('mode', 'explicit_allowlist')
+
+            for level_data in config.get('levels', []):
+                level = ComplianceLevel(
+                    name=level_data['name'],
+                    description=level_data.get('description', ''),
+                    aliases=level_data.get('aliases', []),
+                    allowed_with=level_data.get('allowed_with', [level_data['name']])
+                )
+                self.levels[level.name] = level
+
+                # Map canonical name to itself
+                self._name_to_canonical[level.name] = level.name
+
+                # Map aliases to canonical name
+                for alias in level.aliases:
+                    self._name_to_canonical[alias] = level.name
+
+            logger.info(f"Loaded {len(self.levels)} compliance levels from {config_path}")
+            logger.debug(f"Compliance levels: {list(self.levels.keys())}")
+
+        except Exception as e:
+            logger.error(f"Error loading compliance-levels.json: {e}")
+            # Continue with empty config for permissive validation
+
+    def get_canonical_name(self, name: Optional[str]) -> Optional[str]:
+        """Get the canonical name for a compliance level (resolves aliases).
+
+        Args:
+            name: Compliance level name or alias
+
+        Returns:
+            Canonical name, or None if not found
+        """
+        if not name:
+            return None
+        return self._name_to_canonical.get(name)
+
+    def validate_compliance_level(self, level_name: Optional[str], context: str = "") -> Optional[str]:
+        """Validate a compliance level name.
+
+        Args:
+            level_name: The compliance level to validate
+            context: Context for logging (e.g., "MCP server 'calculator'")
+
+        Returns:
+            Canonical name if valid, None if invalid (with warning logged)
+        """
+        if not level_name:
+            return None
+
+        canonical = self.get_canonical_name(level_name)
+
+        if canonical is None:
+            # No compliance config loaded - permissive mode
+            if not self.levels:
+                return level_name
+
+            # Unknown compliance level
+            valid_levels = list(self.levels.keys())
+            logger.warning(
+                f"Invalid compliance level '{level_name}' {context}. "
+                f"Valid levels: {', '.join(valid_levels)}. "
+                f"Setting to None."
+            )
+            return None
+
+        if canonical != level_name:
+            logger.debug(f"Resolved alias '{level_name}' to '{canonical}' {context}")
+
+        return canonical
+
+    def is_accessible(self, user_level: Optional[str], resource_level: Optional[str]) -> bool:
+        """Check if a resource at resource_level is accessible given user_level.
+
+        In explicit allowlist mode:
+        - Each level defines which other levels can be used together
+        - For example, HIPAA might allow HIPAA and SOC2, but not Public
+        - None (unset) is accessible by all and can access all
+
+        Args:
+            user_level: User's selected compliance level
+            resource_level: Resource's compliance level
+
+        Returns:
+            True if resource is accessible, False otherwise
+        """
+        # If either is None/unset, resource is accessible (backward compatibility)
+        if not user_level or not resource_level:
+            return True
+
+        # Get canonical names
+        user_canonical = self.get_canonical_name(user_level)
+        resource_canonical = self.get_canonical_name(resource_level)
+
+        # If we don't have level info, be permissive
+        if not user_canonical or not resource_canonical:
+            return True
+
+        # Get level object for user
+        user_level_obj = self.levels.get(user_canonical)
+
+        if not user_level_obj:
+            return True
+
+        # Check if resource_level is in the user's allowed_with list
+        return resource_canonical in user_level_obj.allowed_with
+
+    def get_accessible_levels(self, user_level: Optional[str]) -> Set[str]:
+        """Get all compliance levels accessible to a user.
+
+        Args:
+            user_level: User's selected compliance level
+
+        Returns:
+            Set of accessible compliance level names (canonical)
+        """
+        if not user_level or not self.levels:
+            # Return all levels if no user level or no config
+            return set(self.levels.keys()) if self.levels else set()
+
+        user_canonical = self.get_canonical_name(user_level)
+        if not user_canonical or user_canonical not in self.levels:
+            return set(self.levels.keys())
+
+        user_level_obj = self.levels[user_canonical]
+
+        # Return the allowed_with list for this level
+        return set(user_level_obj.allowed_with)
+
+    def get_all_levels(self) -> List[str]:
+        """Get all defined compliance level names (canonical).
+
+        Returns:
+            List of compliance level names in definition order
+        """
+        return list(self.levels.keys())
+
+
+# Global instance
+_compliance_manager: Optional[ComplianceLevelManager] = None
+
+
+def get_compliance_manager() -> ComplianceLevelManager:
+    """Get the global compliance level manager instance."""
+    global _compliance_manager
+    if _compliance_manager is None:
+        _compliance_manager = ComplianceLevelManager()
+    return _compliance_manager