atlas-chat 0.1.0__py3-none-any.whl

atlas/core/domain_whitelist.py

@@ -0,0 +1,147 @@
+"""
+Domain whitelist management for email access control.
+
+Loads domain whitelist definitions from atlas.domain-whitelist.json and provides
+validation for user email domains.
+"""
+
+import json
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Set
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DomainWhitelistConfig:
+    """Configuration for domain whitelist."""
+    domains: Set[str]
+    subdomain_matching: bool
+    version: str
+    description: str
+
+
+class DomainWhitelistManager:
+    """Manages domain whitelist configuration and validation."""
+
+    def __init__(self, config_path: Optional[Path] = None):
+        """Initialize the domain whitelist manager.
+
+        Args:
+            config_path: Path to domain-whitelist.json. If None, uses default location.
+        """
+        self.config: Optional[DomainWhitelistConfig] = None
+        self.config_loaded: bool = False
+
+        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" / "domain-whitelist.json",
+                project_root / "config" / "defaults" / "domain-whitelist.json",
+                backend_root / "configfilesadmin" / "domain-whitelist.json",
+                backend_root / "configfiles" / "domain-whitelist.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 domain-whitelist.json found, whitelist validation disabled")
+            self.config = DomainWhitelistConfig(
+                domains=set(),
+                subdomain_matching=True,
+                version="1.0",
+                description="No config loaded"
+            )
+            self.config_loaded = False
+
+    def _load_config(self, config_path: Path):
+        """Load domain whitelist configuration from JSON file."""
+        try:
+            with open(config_path, 'r', encoding='utf-8') as f:
+                config_data = json.load(f)
+
+            # Extract domains from the list of domain objects
+            domains = set()
+            for domain_entry in config_data.get('domains', []):
+                if isinstance(domain_entry, dict):
+                    domains.add(domain_entry.get('domain', '').lower())
+                elif isinstance(domain_entry, str):
+                    domains.add(domain_entry.lower())
+
+            self.config = DomainWhitelistConfig(
+                domains=domains,
+                subdomain_matching=config_data.get('subdomain_matching', True),
+                version=config_data.get('version', '1.0'),
+                description=config_data.get('description', '')
+            )
+            self.config_loaded = True
+
+            logger.info(f"Loaded {len(self.config.domains)} domains from {config_path}")
+
+        except Exception as e:
+            logger.error(f"Error loading domain-whitelist.json: {e}")
+            logger.warning("Whitelist validation disabled due to config error")
+            # Use empty config on error
+            self.config = DomainWhitelistConfig(
+                domains=set(),
+                subdomain_matching=True,
+                version="1.0",
+                description="Error loading config"
+            )
+            self.config_loaded = False
+
+    def is_domain_allowed(self, email: str) -> bool:
+        """Check if an email address is from an allowed domain.
+
+        Note: This method only validates against the whitelist.
+        The FEATURE_DOMAIN_WHITELIST_ENABLED flag controls whether
+        the middleware uses this validation.
+
+        Args:
+            email: Email address to validate
+
+        Returns:
+            True if domain is allowed, False otherwise
+        """
+        # If config wasn't successfully loaded, allow all (fail open)
+        if not self.config_loaded:
+            return True
+
+        if not email or "@" not in email:
+            return False
+
+        domain = email.split("@", 1)[1].lower()
+
+        # Check if domain is in whitelist (O(1) lookup)
+        if domain in self.config.domains:
+            return True
+
+        # Check subdomains if enabled - check each parent level
+        if self.config.subdomain_matching:
+            # Split domain and check each parent level
+            # e.g., for "mail.dept.sandia.gov" check: "dept.sandia.gov", "sandia.gov"
+            parts = domain.split(".")
+            for i in range(1, len(parts)):
+                parent_domain = ".".join(parts[i:])
+                if parent_domain in self.config.domains:
+                    return True
+
+        return False
+
+    def get_domains(self) -> Set[str]:
+        """Get the set of whitelisted domains.
+
+        Returns:
+            Set of allowed domains
+        """
+        return self.config.domains if self.config else set()

atlas/core/domain_whitelist_middleware.py

@@ -0,0 +1,82 @@
+"""Email domain whitelist validation middleware.
+
+This middleware enforces that users must have email addresses from whitelisted
+domains. Enabled/disabled via the FEATURE_DOMAIN_WHITELIST_ENABLED feature flag.
+Domain list is loaded from atlas.domain-whitelist.json.
+"""
+
+import logging
+
+from fastapi import Request
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import JSONResponse, RedirectResponse, Response
+
+from atlas.core.domain_whitelist import DomainWhitelistManager
+
+logger = logging.getLogger(__name__)
+
+
+class DomainWhitelistMiddleware(BaseHTTPMiddleware):
+    """Middleware to enforce email domain whitelist restrictions."""
+
+    def __init__(self, app, auth_redirect_url: str = "/auth"):
+        """Initialize domain whitelist middleware.
+
+        Args:
+            app: ASGI application
+            auth_redirect_url: URL to redirect to on auth failure (default: /auth)
+        """
+        super().__init__(app)
+        self.auth_redirect_url = auth_redirect_url
+        self.whitelist_manager = DomainWhitelistManager()
+
+        logger.info(f"Domain whitelist middleware loaded: {len(self.whitelist_manager.get_domains())} domains (config_loaded={self.whitelist_manager.config_loaded})")
+
+    async def dispatch(self, request: Request, call_next) -> Response:
+        """Check if user email is from a whitelisted domain.
+
+        Args:
+            request: Incoming HTTP request
+            call_next: Next middleware/handler in chain
+
+        Returns:
+            Response from next handler if authorized, or 403/redirect if not
+        """
+        # Skip check for health endpoint and auth redirect endpoint
+        if request.url.path == '/api/health' or request.url.path == self.auth_redirect_url:
+            return await call_next(request)
+
+        # Get email from request state (set by AuthMiddleware)
+        email = getattr(request.state, "user_email", None)
+
+        if not email or "@" not in email:
+            logger.warning("Domain whitelist check failed: missing or invalid email")
+            return self._unauthorized_response(request, "User email required")
+
+        # Check if domain is allowed
+        if not self.whitelist_manager.is_domain_allowed(email):
+            domain = email.split("@", 1)[1].lower()
+            logger.warning(f"Domain whitelist check failed: unauthorized domain {domain}")
+            return self._unauthorized_response(
+                request,
+                "Access restricted to whitelisted domains"
+            )
+
+        return await call_next(request)
+
+    def _unauthorized_response(self, request: Request, detail: str) -> Response:
+        """Return appropriate unauthorized response based on endpoint type.
+
+        Args:
+            request: Incoming HTTP request
+            detail: Error detail message
+
+        Returns:
+            JSONResponse for API endpoints, RedirectResponse for others
+        """
+        if request.url.path.startswith('/api/'):
+            return JSONResponse(
+                status_code=403,
+                content={"detail": detail}
+            )
+        return RedirectResponse(url=self.auth_redirect_url, status_code=302)

atlas/core/http_client.py

@@ -0,0 +1,28 @@
+"""
+Minimal HTTP client stub for basic chat functionality.
+This is a temporary implementation for testing.
+"""
+
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def create_rag_client(base_url: str = "", timeout: float = 30.0) -> Any:
+    """
+    Create a simple RAG client stub.
+    For basic chat, this just returns a mock client.
+    """
+    class MockRAGClient:
+        def __init__(self):
+            pass
+
+        async def query(self, *args, **kwargs):
+            """Mock RAG query - returns empty result."""
+            return {
+                "content": "RAG not available in basic chat mode",
+                "metadata": {}
+            }
+
+    return MockRAGClient()

atlas/core/log_sanitizer.py

@@ -0,0 +1,102 @@
+"""
+Minimal utilities for basic chat functionality.
+"""
+
+import logging
+import re
+from typing import Any
+
+from fastapi import Request
+
+logger = logging.getLogger(__name__)
+
+_CONTROL_CHARS_RE = re.compile(r'[\x00-\x1f\x7f-\x9f]')
+# Matches Unicode line separators (LINE SEPARATOR and PARAGRAPH SEPARATOR)
+_UNICODE_NEWLINES_RE = re.compile(r'[\u2028\u2029]')
+# Matches explicit CR, LF, and CRLF for maximal coverage
+_STANDARD_NEWLINES_RE = re.compile(r'(\r\n|\r|\n)')
+
+def sanitize_for_logging(value: Any) -> str:
+    """
+    Sanitize a value for safe logging by removing ALL newlines (including Unicode and CRLF)
+    and control characters, to defend against log injection.
+
+    Removes ASCII control characters (C0 and C1 ranges), CR/LF in any combination,
+    and Unicode line/paragraph separators. This includes characters
+    like newlines (\\n, \\r, \\r\\n, U+2028, U+2029), tabs, escape sequences, and other
+    non-printable characters that could be used to manipulate log output or inject fake log entries.
+
+    Args:
+        value: Any value to sanitize. If not a string, it will be converted
+               to string representation first.
+
+    Returns:
+        str: Sanitized string with all control and newline characters removed.
+
+    Examples:
+        >>> sanitize_for_logging("Hello\\nWorld")
+        'HelloWorld'
+        >>> sanitize_for_logging("Test\\x1b[31mRed\\x1b[0m")
+        'TestRed'
+        >>> sanitize_for_logging("Fake\u2028Log")
+        'FakeLog'
+        >>> sanitize_for_logging("line1\\r\\nline2\\rline3\\nline4")
+        'line1line2line3line4'
+        >>> sanitize_for_logging("A\u2028B\u2029C")
+        'ABC'
+        >>> sanitize_for_logging(123)
+        '123'
+    """
+    if value is None:
+        return ''
+    if not isinstance(value, str):
+        value = str(value)
+    value = _CONTROL_CHARS_RE.sub('', value)
+    value = _UNICODE_NEWLINES_RE.sub('', value)
+    value = _STANDARD_NEWLINES_RE.sub('', value)
+    return value
+
+
+def summarize_tool_approval_response_for_logging(data: Any) -> str:
+    """Return a non-sensitive summary of a tool approval response payload.
+
+    This is intentionally conservative: it never logs tool argument values or
+    rejection reasons because these can contain sensitive user content.
+
+    Expected input shape (from websocket):
+        {
+          "type": "tool_approval_response",
+          "tool_call_id": "...",
+          "approved": true/false,
+          "arguments": {...},
+          "reason": "..."
+        }
+    """
+    if not isinstance(data, dict):
+        return f"type=tool_approval_response payload_type={sanitize_for_logging(type(data).__name__)}"
+
+    tool_call_id = sanitize_for_logging(data.get("tool_call_id"))
+    approved_raw = data.get("approved", False)
+    approved = bool(approved_raw)
+
+    arguments = data.get("arguments")
+    has_arguments = arguments is not None
+    arguments_count = len(arguments) if isinstance(arguments, dict) else (1 if has_arguments else 0)
+
+    reason = data.get("reason")
+    has_reason = bool(reason)
+
+    return (
+        "type=tool_approval_response "
+        f"tool_call_id={tool_call_id} "
+        f"approved={approved} "
+        f"has_arguments={has_arguments} "
+        f"arguments_count={arguments_count} "
+        f"has_reason={has_reason}"
+    )
+
+
+
+async def get_current_user(request: Request) -> str:
+    """Get current user from request state (set by middleware)."""
+    return getattr(request.state, 'user_email', 'test@test.com')

atlas/core/metrics_logger.py

@@ -0,0 +1,59 @@
+"""
+Metrics logging utility for tracking user activities without capturing sensitive data.
+
+This module provides a centralized way to log user activity metrics that:
+- Use the [METRIC] prefix for easy filtering
+- Include the username for tracking
+- Only log metadata (counts, sizes, types)
+- NEVER log sensitive data like prompts, tool arguments, filenames, or error details
+
+Usage:
+    from atlas.core.metrics_logger import log_metric
+
+    log_metric("llm_call", user_email, model="gpt-4", message_count=5)
+    log_metric("tool_call", user_email, tool_name="calculator")
+    log_metric("file_upload", user_email, file_size=1024, content_type="application/pdf")
+    log_metric("error", user_email, error_type="rate_limit")
+"""
+
+import logging
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+def log_metric(
+    event_type: str,
+    user_email: Optional[str] = None,
+    **kwargs: Any
+) -> None:
+    """
+    Log a metric event for user activity tracking.
+
+    This function respects the FEATURE_METRICS_LOGGING_ENABLED setting.
+    When disabled, no metrics are logged.
+
+    Args:
+        event_type: Type of event (e.g., "llm_call", "tool_call", "file_upload", "error")
+        user_email: User's email address (will be sanitized)
+        **kwargs: Additional metadata to log (only non-sensitive data)
+    """
+    # Import here to avoid circular dependencies
+    from atlas.core.log_sanitizer import sanitize_for_logging
+    from atlas.modules.config import config_manager
+
+    if not config_manager.app_settings.feature_metrics_logging_enabled:
+        return
+
+    sanitized_user = sanitize_for_logging(user_email) if user_email else "unknown"
+
+    parts = [f"[METRIC] [{sanitized_user}] {event_type}"]
+
+    if kwargs:
+        metadata_parts = [
+            f"{key}={sanitize_for_logging(value)}"
+            for key, value in kwargs.items()
+        ]
+        parts.append(" ".join(metadata_parts))
+
+    logger.info(" ".join(parts))

atlas/core/middleware.py

@@ -0,0 +1,131 @@
+"""FastAPI middleware for authentication and logging."""
+
+import logging
+
+from fastapi import Request
+from fastapi.responses import JSONResponse, RedirectResponse
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import Response
+
+from atlas.core.auth import get_user_from_aws_alb_jwt, get_user_from_header
+from atlas.core.capabilities import verify_file_token
+from atlas.infrastructure.app_factory import app_factory
+
+logger = logging.getLogger(__name__)
+
+
+class AuthMiddleware(BaseHTTPMiddleware):
+    """Middleware to handle authentication and logging."""
+
+    def __init__(
+        self,
+        app,
+        debug_mode: bool = False,
+        auth_header_name: str = "X-User-Email",
+        auth_header_type: str = "email-string",
+        auth_aws_expected_alb_arn: str = "",
+        auth_aws_region: str = "us-east-1",
+        proxy_secret_enabled: bool = False,
+        proxy_secret_header: str = "X-Proxy-Secret",
+        proxy_secret: str = None,
+        auth_redirect_url: str = "/auth"
+    ):
+        super().__init__(app)
+        self.debug_mode = debug_mode
+        self.auth_header_name = auth_header_name
+        self.auth_header_type = auth_header_type
+        self.auth_aws_expected_alb_arn = auth_aws_expected_alb_arn
+        self.auth_aws_region = auth_aws_region
+        self.proxy_secret_enabled = proxy_secret_enabled
+        self.proxy_secret_header = proxy_secret_header
+        self.proxy_secret = proxy_secret
+        self.auth_redirect_url = auth_redirect_url
+
+    async def dispatch(self, request: Request, call_next) -> Response:
+        # Log request
+        logger.debug("Request: %s %s", request.method, request.url.path)
+
+        # Skip auth for static files, health check, and configured auth endpoint
+        if (request.url.path.startswith('/static') or
+            request.url.path == '/api/health' or
+            request.url.path == self.auth_redirect_url):
+            return await call_next(request)
+
+        # Validate proxy secret if enabled (skip in debug mode for local development)
+        if self.proxy_secret_enabled and self.proxy_secret and not self.debug_mode:
+            proxy_secret_value = request.headers.get(self.proxy_secret_header)
+
+            if not proxy_secret_value or proxy_secret_value != self.proxy_secret:
+                logger.warning(f"Invalid or missing proxy secret for {request.url.path}")
+                # Distinguish between API endpoints (return 401) and browser endpoints (redirect)
+                if request.url.path.startswith('/api/'):
+                    return JSONResponse(
+                        status_code=401,
+                        content={"detail": "Unauthorized: Invalid proxy secret"}
+                    )
+                else:
+                    return RedirectResponse(url=self.auth_redirect_url, status_code=302)
+
+        # Check for capability token in download URLs (allows MCP servers to access files)
+        if request.url.path.startswith('/api/files/download/'):
+            token = request.query_params.get('token')
+            if token:
+                claims = verify_file_token(token)
+                if claims:
+                    # Valid capability token - extract user from token and allow request
+                    # Note: We only validate token authenticity here (authentication).
+                    # The route handler validates that token's file key matches the requested
+                    # file (authorization). This separation of concerns keeps middleware focused
+                    # on authentication while route handlers handle resource-specific authorization.
+                    user_email = claims.get('u')
+                    if user_email:
+                        logger.debug("Authenticated via capability token for user: %s", user_email)
+                        request.state.user_email = user_email
+                        return await call_next(request)
+                    else:
+                        logger.warning("Valid token but missing user email claim")
+                else:
+                    logger.warning("Invalid capability token provided")
+
+        # Check authentication via configured header (default: X-User-Email)
+        user_email = None
+        if self.debug_mode:
+            # In debug mode, honor auth header if provided, otherwise use config test user
+            x_auth_header = request.headers.get(self.auth_header_name)
+            if x_auth_header:
+                # Apply same authentication logic as production for testing
+                if self.auth_header_type == "aws-alb-jwt":
+                    user_email = get_user_from_aws_alb_jwt(x_auth_header, self.auth_aws_expected_alb_arn, self.auth_aws_region)
+                else:
+                    user_email = get_user_from_header(x_auth_header)
+            else:
+                # Get test user from config
+                config_manager = app_factory.get_config_manager()
+                user_email = config_manager.app_settings.test_user
+            # logger.info(f"Debug mode: using user {user_email}")
+        else:
+            x_auth_header = request.headers.get(self.auth_header_name)
+
+            # Extract the user's email, depending on the datatype of auth header
+            if self.auth_header_type == "aws-alb-jwt": # Amazon Application Load Balancer
+                user_email = get_user_from_aws_alb_jwt(x_auth_header, self.auth_aws_expected_alb_arn, self.auth_aws_region)
+            else:
+                user_email = get_user_from_header(x_auth_header)
+
+            if not user_email:
+                # Distinguish between API endpoints (return 401) and browser endpoints (redirect)
+                if request.url.path.startswith('/api/'):
+                    logger.warning(f"Missing authentication for API endpoint: {request.url.path}")
+                    return JSONResponse(
+                        status_code=401,
+                        content={"detail": "Unauthorized"}
+                    )
+                else:
+                    logger.warning(f"Missing {self.auth_header_name}, redirecting to {self.auth_redirect_url}")
+                    return RedirectResponse(url=self.auth_redirect_url, status_code=302)
+
+        # Add user to request state
+        request.state.user_email = user_email
+
+        response = await call_next(request)
+        return response