mdb-engine 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

mdb_engine/auth/websocket_sessions.py

@@ -0,0 +1,433 @@
+"""
+WebSocket Session Manager with Envelope Encryption
+
+Manages WebSocket session keys using envelope encryption and private collections.
+Provides secure-by-default WebSocket authentication without relying on CSRF cookies.
+
+This module is part of MDB_ENGINE - MongoDB Engine.
+
+Security Model:
+- Session keys generated on authentication
+- Stored encrypted in _mdb_engine_websocket_sessions collection
+- Validated during WebSocket upgrade
+- Uses envelope encryption (same as app secrets)
+- Security by default: CSRF always required
+"""
+
+import base64
+import logging
+import secrets
+from collections.abc import Callable
+from datetime import datetime, timedelta
+from typing import Any
+
+from motor.motor_asyncio import AsyncIOMotorDatabase
+from pymongo.errors import OperationFailure, PyMongoError
+
+from ..core.encryption import EnvelopeEncryptionService
+
+logger = logging.getLogger(__name__)
+
+# Collection name for storing encrypted WebSocket session keys
+WEBSOCKET_SESSIONS_COLLECTION_NAME = "_mdb_engine_websocket_sessions"
+
+# Session key configuration
+SESSION_KEY_SIZE = 32  # 256 bits
+SESSION_TTL_HOURS = 24  # Sessions expire after 24 hours
+
+
+class WebSocketSessionManager:
+    """
+    Manages WebSocket session keys using envelope encryption.
+
+    Session keys are:
+    - Generated on user authentication
+    - Encrypted using envelope encryption
+    - Stored in private collection (_mdb_engine_websocket_sessions)
+    - Validated during WebSocket upgrade
+    - Automatically expired after TTL
+    """
+
+    def __init__(
+        self,
+        mongo_db: AsyncIOMotorDatabase,
+        encryption_service: EnvelopeEncryptionService,
+    ):
+        """
+        Initialize the WebSocket session manager.
+
+        Args:
+            mongo_db: MongoDB database instance (raw, not scoped)
+            encryption_service: Envelope encryption service instance
+        """
+        self._mongo_db = mongo_db
+        self._encryption_service = encryption_service
+        self._sessions_collection = mongo_db[WEBSOCKET_SESSIONS_COLLECTION_NAME]
+
+    @staticmethod
+    def generate_session_key() -> str:
+        """
+        Generate a random WebSocket session key.
+
+        Returns:
+            Base64-encoded session key string
+        """
+        key_bytes = secrets.token_bytes(SESSION_KEY_SIZE)
+        return base64.urlsafe_b64encode(key_bytes).decode().rstrip("=")
+
+    async def create_session(
+        self,
+        user_id: str,
+        user_email: str | None = None,
+        app_slug: str | None = None,
+    ) -> str:
+        """
+        Create a new WebSocket session with encrypted session key.
+
+        Args:
+            user_id: User ID
+            user_email: Optional user email
+            app_slug: Optional app slug for scoping
+
+        Returns:
+            Plaintext session key (to be sent to client)
+
+        Raises:
+            OperationFailure: If MongoDB operation fails
+        """
+        try:
+            # Generate session key
+            session_key = self.generate_session_key()
+
+            # Encrypt session key using envelope encryption
+            encrypted_key, encrypted_dek = self._encryption_service.encrypt_secret(session_key)
+
+            # Encode as base64 for storage
+            encrypted_key_b64 = base64.b64encode(encrypted_key).decode()
+            encrypted_dek_b64 = base64.b64encode(encrypted_dek).decode()
+
+            # Calculate expiration
+            expires_at = datetime.utcnow() + timedelta(hours=SESSION_TTL_HOURS)
+
+            # Prepare document
+            document = {
+                "_id": session_key,  # Use session key as ID for fast lookup
+                "user_id": user_id,
+                "user_email": user_email,
+                "app_slug": app_slug,
+                "encrypted_key": encrypted_key_b64,
+                "encrypted_dek": encrypted_dek_b64,
+                "algorithm": "AES-256-GCM",
+                "created_at": datetime.utcnow(),
+                "expires_at": expires_at,
+            }
+
+            # Store in private collection
+            await self._sessions_collection.insert_one(document)
+
+            logger.info(
+                f"Created WebSocket session for user '{user_id}' "
+                f"(app: {app_slug}, expires: {expires_at})"
+            )
+
+            return session_key
+
+        except (OperationFailure, PyMongoError):
+            logger.exception("Failed to create WebSocket session")
+            raise
+
+    async def validate_session(
+        self,
+        session_key: str,
+        user_id: str | None = None,
+    ) -> dict[str, Any] | None:
+        """
+        Validate a WebSocket session key.
+
+        Args:
+            session_key: Session key to validate
+            user_id: Optional user ID for additional validation
+
+        Returns:
+            Session document if valid, None otherwise
+
+        Raises:
+            OperationFailure: If MongoDB operation fails
+        """
+        try:
+            # Find session by key
+            session_doc = await self._sessions_collection.find_one({"_id": session_key})
+
+            if not session_doc:
+                logger.warning(f"WebSocket session not found: {session_key[:16]}...")
+                return None
+
+            # Check expiration
+            expires_at = session_doc.get("expires_at")
+            if expires_at and expires_at < datetime.utcnow():
+                logger.warning(
+                    f"WebSocket session expired: {session_key[:16]}... " f"(expired: {expires_at})"
+                )
+                # Clean up expired session
+                await self._sessions_collection.delete_one({"_id": session_key})
+                return None
+
+            # Optional: Validate user_id matches
+            if user_id and session_doc.get("user_id") != user_id:
+                logger.warning(
+                    f"WebSocket session user mismatch: "
+                    f"session_user={session_doc.get('user_id')}, "
+                    f"provided_user={user_id}"
+                )
+                return None
+
+            # Decrypt session key to verify it's valid
+            try:
+                encrypted_key = base64.b64decode(session_doc["encrypted_key"])
+                encrypted_dek = base64.b64decode(session_doc["encrypted_dek"])
+                decrypted_key = self._encryption_service.decrypt_secret(
+                    encrypted_key, encrypted_dek
+                )
+
+                # Verify decrypted key matches session_key
+                if decrypted_key != session_key:
+                    logger.error(
+                        f"WebSocket session key decryption mismatch: "
+                        f"session_key={session_key[:16]}..."
+                    )
+                    return None
+
+            except (ValueError, TypeError, AttributeError, KeyError):
+                logger.exception("Failed to decrypt WebSocket session key")
+                return None
+
+            logger.debug(
+                f"Validated WebSocket session for user '{session_doc.get('user_id')}' "
+                f"(app: {session_doc.get('app_slug')})"
+            )
+
+            return {
+                "user_id": session_doc.get("user_id"),
+                "user_email": session_doc.get("user_email"),
+                "app_slug": session_doc.get("app_slug"),
+                "created_at": session_doc.get("created_at"),
+                "expires_at": session_doc.get("expires_at"),
+            }
+
+        except (OperationFailure, PyMongoError):
+            logger.exception("Failed to validate WebSocket session")
+            raise
+
+    async def revoke_session(self, session_key: str) -> bool:
+        """
+        Revoke a WebSocket session.
+
+        Args:
+            session_key: Session key to revoke
+
+        Returns:
+            True if session was revoked, False if not found
+        """
+        try:
+            result = await self._sessions_collection.delete_one({"_id": session_key})
+            if result.deleted_count > 0:
+                logger.info(f"Revoked WebSocket session: {session_key[:16]}...")
+                return True
+            return False
+        except (OperationFailure, PyMongoError):
+            logger.exception("Failed to revoke WebSocket session")
+            return False
+
+    async def revoke_user_sessions(self, user_id: str, app_slug: str | None = None) -> int:
+        """
+        Revoke all sessions for a user.
+
+        Args:
+            user_id: User ID
+            app_slug: Optional app slug filter
+
+        Returns:
+            Number of sessions revoked
+        """
+        try:
+            query = {"user_id": user_id}
+            if app_slug:
+                query["app_slug"] = app_slug
+
+            result = await self._sessions_collection.delete_many(query)
+            logger.info(
+                f"Revoked {result.deleted_count} WebSocket sessions "
+                f"for user '{user_id}' (app: {app_slug})"
+            )
+            return result.deleted_count
+        except (OperationFailure, PyMongoError):
+            logger.exception("Failed to revoke user WebSocket sessions")
+            return 0
+
+    async def cleanup_expired_sessions(self) -> int:
+        """
+        Clean up expired WebSocket sessions.
+
+        Returns:
+            Number of sessions cleaned up
+        """
+        try:
+            result = await self._sessions_collection.delete_many(
+                {"expires_at": {"$lt": datetime.utcnow()}}
+            )
+            if result.deleted_count > 0:
+                logger.info(f"Cleaned up {result.deleted_count} expired WebSocket sessions")
+            return result.deleted_count
+        except (OperationFailure, PyMongoError):
+            logger.exception("Failed to cleanup expired WebSocket sessions")
+            return 0
+
+
+def create_websocket_session_endpoint(
+    session_manager: WebSocketSessionManager,
+) -> Callable:
+    """
+    Create a FastAPI endpoint for generating WebSocket session keys.
+
+    This endpoint requires authentication and generates a new WebSocket session key
+    for the authenticated user. The session key is encrypted and stored in the
+    private collection.
+
+    Args:
+        session_manager: WebSocketSessionManager instance
+
+    Returns:
+        FastAPI route handler function
+
+    Example:
+        ```python
+        from mdb_engine.auth.websocket_sessions import (
+            WebSocketSessionManager,
+            create_websocket_session_endpoint,
+        )
+        from mdb_engine.core.encryption import EnvelopeEncryptionService
+
+        # Initialize session manager
+        encryption_service = EnvelopeEncryptionService()
+        session_manager = WebSocketSessionManager(
+            mongo_db=db,
+            encryption_service=encryption_service,
+        )
+
+        # Create endpoint
+        endpoint = create_websocket_session_endpoint(session_manager)
+        app.get("/auth/websocket-session")(endpoint)
+        ```
+
+    The endpoint:
+    - Requires authentication (user must be logged in)
+    - Returns JSON: `{"session_key": "...", "expires_at": "..."}`
+    - Uses user info from `request.state.user` (set by SharedAuthMiddleware)
+    """
+    from fastapi import Request, status
+    from fastapi.responses import JSONResponse
+
+    async def websocket_session_endpoint(request: Request) -> JSONResponse:
+        """
+        Generate a WebSocket session key for the authenticated user.
+
+        Requires:
+        - User to be authenticated (via request.state.user or auth cookie)
+        - WebSocket session manager to be available
+
+        Returns:
+        - JSONResponse with session_key and expires_at
+        """
+        # Check if user is authenticated (set by middleware)
+        user = getattr(request.state, "user", None)
+
+        # If not set by middleware, try to authenticate using cookie
+        # This handles the case where endpoint is on parent app without auth middleware
+        if not user:
+            from .shared_middleware import AUTH_COOKIE_NAME
+
+            # Get user pool from app state
+            user_pool = None
+            try:
+                if hasattr(request, "app") and hasattr(request.app, "state"):
+                    user_pool = getattr(request.app.state, "user_pool", None)
+            except (AttributeError, TypeError):
+                pass
+
+            # Only try to authenticate if we have a real user pool (not None)
+            if user_pool is not None:
+                # Extract token from cookie
+                token = None
+                try:
+                    if hasattr(request, "cookies"):
+                        token = request.cookies.get(AUTH_COOKIE_NAME)
+                except (AttributeError, TypeError):
+                    pass
+
+                if token:
+                    try:
+                        # Validate token and get user
+                        user = await user_pool.validate_token(token)
+                    except (TypeError, AttributeError):
+                        # If user_pool is a mock that can't be awaited, ignore
+                        pass
+
+            if not user:
+                return JSONResponse(
+                    status_code=status.HTTP_401_UNAUTHORIZED,
+                    content={"detail": "Authentication required"},
+                )
+
+        # Extract user info
+        # Prefer user_id, sub (JWT standard), or _id (MongoDB document ID)
+        user_id = user.get("user_id") or user.get("sub") or user.get("_id")
+        if not user_id:
+            # Email is not a valid user_id - it's just metadata
+            logger.error("Cannot generate WebSocket session: user_id not found in user data")
+            return JSONResponse(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                content={"detail": "Invalid user data"},
+            )
+        user_email = user.get("email")
+        app_slug = getattr(request.state, "app_slug", None)
+
+        try:
+            # Generate session key
+            session_key = await session_manager.create_session(
+                user_id=str(user_id),
+                user_email=user_email,
+                app_slug=app_slug,
+            )
+
+            # Get expiration time (24 hours from now)
+            from datetime import datetime, timedelta
+
+            expires_at = datetime.utcnow() + timedelta(hours=SESSION_TTL_HOURS)
+
+            logger.info(
+                f"Generated WebSocket session key for user '{user_id}' " f"(app: {app_slug})"
+            )
+
+            return JSONResponse(
+                {
+                    "session_key": session_key,
+                    "expires_at": expires_at.isoformat(),
+                    "ttl_hours": SESSION_TTL_HOURS,
+                }
+            )
+
+        except (
+            ValueError,
+            TypeError,
+            AttributeError,
+            RuntimeError,
+            OperationFailure,
+            PyMongoError,
+        ):
+            logger.exception("Failed to generate WebSocket session key")
+            return JSONResponse(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                content={"detail": "Failed to generate WebSocket session key"},
+            )
+
+    return websocket_session_endpoint

mdb_engine/core/engine.py

@@ -148,6 +148,7 @@ class MongoDBEngine:
         self._service_initializer: ServiceInitializer | None = None
         self._encryption_service: EnvelopeEncryptionService | None = None
         self._app_secrets_manager: AppSecretsManager | None = None
+        self._websocket_session_manager: Any | None = None  # WebSocketSessionManager
 
         # Store app read_scopes mapping for validation
         self._app_read_scopes: dict[str, list[str]] = {}
@@ -201,6 +202,13 @@ class MongoDBEngine:
                 mongo_db=self._connection_manager.mongo_db,
                 encryption_service=self._encryption_service,
             )
+            # Initialize WebSocket session manager for secure-by-default WebSocket auth
+            from ..auth.websocket_sessions import WebSocketSessionManager
+
+            self._websocket_session_manager = WebSocketSessionManager(
+                mongo_db=self._connection_manager.mongo_db,
+                encryption_service=self._encryption_service,
+            )
 
         # Set up component managers
         self._app_registration_manager = AppRegistrationManager(
@@ -2283,6 +2291,11 @@ class MongoDBEngine:
                 logger.debug(f"No WebSocket configuration found for app '{slug}'")
                 return
 
+            # Store WebSocket config in parent app state for CSRF middleware to access
+            if not hasattr(parent_app.state, "websocket_configs"):
+                parent_app.state.websocket_configs = {}
+            parent_app.state.websocket_configs[slug] = websockets_config
+
             try:
                 from fastapi import APIRouter
 
@@ -2489,6 +2502,13 @@ class MongoDBEngine:
                             child_app.state.audit_log = app.state.audit_log
                         logger.debug(f"Shared user_pool with child app '{slug}'")
 
+                    # Share WebSocket session manager with child app
+                    if hasattr(app.state, "websocket_session_manager"):
+                        child_app.state.websocket_session_manager = (
+                            app.state.websocket_session_manager
+                        )
+                        logger.debug(f"Shared WebSocket session manager with child app '{slug}'")
+
                     # Add middleware for app context helpers
                     from starlette.middleware.base import BaseHTTPMiddleware
                     from starlette.requests import Request
@@ -2836,12 +2856,31 @@ class MongoDBEngine:
             # Create CSRF middleware with default config (will use parent app's CORS config)
             # Exempt routes that don't need CSRF (health checks, public routes from child apps)
             # all_public_routes includes base routes + child app public routes with path prefixes
+            # Add WebSocket session endpoint to public routes (it handles its own auth)
+            public_routes_with_session_endpoint = list(all_public_routes) + [
+                "/auth/websocket-session"
+            ]
             parent_csrf_config = {
                 "csrf_protection": True,
-                "public_routes": all_public_routes,
+                "public_routes": public_routes_with_session_endpoint,
             }
             csrf_middleware = create_csrf_middleware(parent_csrf_config)
             parent_app.add_middleware(csrf_middleware)
+
+            # Store WebSocket session manager in app state for CSRF middleware and endpoints
+            if self._websocket_session_manager:
+                parent_app.state.websocket_session_manager = self._websocket_session_manager
+                logger.info("WebSocket session manager stored in parent app state")
+
+                # Register WebSocket session endpoint on parent app
+                from ..auth.websocket_sessions import create_websocket_session_endpoint
+
+                session_endpoint = create_websocket_session_endpoint(
+                    self._websocket_session_manager
+                )
+                parent_app.get("/auth/websocket-session")(session_endpoint)
+                logger.info("WebSocket session endpoint registered at /auth/websocket-session")
+
             logger.info("CSRFMiddleware added to parent app for WebSocket origin validation")
 
         # Add shared CORS middleware if configured
@@ -3320,6 +3359,7 @@ class MongoDBEngine:
                 self._shared_user_pool = SharedUserPool(
                     self._connection_manager.mongo_db,
                     allow_insecure_dev=is_dev,
+                    websocket_session_manager=self._websocket_session_manager,
                 )
                 await self._shared_user_pool.ensure_indexes()
                 logger.info("SharedUserPool initialized")

mdb_engine/core/manifest.py

@@ -1338,6 +1338,18 @@ MANIFEST_SCHEMA_V2 = {
                                         "auth is required (default: false)"
                                     ),
                                 },
+                                "csrf_required": {
+                                    "type": "boolean",
+                                    "default": True,
+                                    "description": (
+                                        "Require CSRF validation for WebSocket connections "
+                                        "(default: true - security by default). "
+                                        "When true, uses encrypted session keys stored in "
+                                        "private collection for CSRF protection. "
+                                        "Set to false to use Origin validation + "
+                                        "SameSite cookies only."
+                                    ),
+                                },
                             },
                             "additionalProperties": False,
                             "description": (

mdb_engine/core/types.py

@@ -243,6 +243,7 @@ class WebSocketAuthDict(TypedDict, total=False):
 
     required: bool
     allow_anonymous: bool
+    csrf_required: bool  # Whether CSRF cookie is required (default: False)
 
 
 class WebSocketEndpointDict(TypedDict, total=False):

mdb_engine/routing/websockets.py

@@ -365,12 +365,11 @@ async def authenticate_websocket(
     require_auth: bool = True,
 ) -> tuple[str | None, str | None]:
     """
-    Authenticate a WebSocket connection via httpOnly cookies.
+    Authenticate a WebSocket connection via session key or httpOnly cookies.
 
-    Uses cookie-based authentication with CSRF protection:
-    - Token stored in httpOnly cookie (not accessible to JavaScript)
-    - CSRF token validated via double-submit cookie pattern
-    - Origin validation provides additional protection
+    Authentication methods (in order of preference):
+    1. Session key (query param or header) - secure-by-default, uses envelope encryption
+    2. Cookie-based authentication - backward compatibility fallback
 
     Args:
         websocket: FastAPI WebSocket instance (can access headers before accept)
@@ -394,22 +393,71 @@ async def authenticate_websocket(
         return None, None
 
     try:
-        # Extract token from httpOnly cookie
+        # Try to get WebSocket session manager from app
+        websocket_session_manager = None
+        try:
+            app = getattr(websocket, "app", None)
+            if app:
+                websocket_session_manager = getattr(app.state, "websocket_session_manager", None)
+        except (AttributeError, TypeError):
+            pass
+
+        # Method 1: Try session key authentication (secure-by-default)
+        session_key = None
+        try:
+            # Check query params first
+            if hasattr(websocket, "query_params"):
+                session_key = websocket.query_params.get("session_key")
+
+            # Check headers if not in query params
+            if not session_key and hasattr(websocket, "headers"):
+                session_key = websocket.headers.get("X-WebSocket-Session-Key")
+        except (AttributeError, TypeError, KeyError):
+            pass
+
+        if session_key and websocket_session_manager:
+            try:
+                # Validate session key
+                session_data = await websocket_session_manager.validate_session(session_key)
+                if session_data:
+                    user_id = session_data.get("user_id")
+                    user_email = session_data.get("user_email")
+
+                    logger.info(
+                        f"WebSocket authenticated successfully for app '{app_slug}': {user_email} "
+                        f"(method: session_key)"
+                    )
+                    return user_id, user_email
+                else:
+                    logger.warning(
+                        f"WebSocket session key validation failed for app '{app_slug}'. "
+                        f"Session key: {session_key[:16]}..."
+                    )
+            except (ValueError, TypeError, AttributeError, KeyError, RuntimeError) as e:
+                logger.warning(f"WebSocket session key validation error for app '{app_slug}': {e}")
+                # Fall through to cookie-based auth
+
+        # Method 2: Fall back to cookie-based authentication (backward compatibility)
+        from ..auth.shared_middleware import AUTH_COOKIE_NAME
+
         cookies = _get_cookies_from_websocket(websocket)
-        token = cookies.get("token")  # Standard auth token cookie name
+        token = cookies.get(AUTH_COOKIE_NAME)  # Use mdb_auth_token (same as shared middleware)
 
         if not token:
-            logger.warning(
-                f"No token cookie found for WebSocket connection to app '{app_slug}' "
+            logger.error(
+                f"❌ No authentication found for WebSocket connection to app '{app_slug}' "
                 f"(require_auth={require_auth}). "
-                f"Ensure httpOnly cookie is set during authentication."
+                f"Session key: {bool(session_key)}, Cookie: {bool(token)}, "
+                f"Available cookies: {list(cookies.keys()) if cookies else 'none'}. "
+                f"Ensure session key or httpOnly cookie is set during authentication."
             )
             if require_auth:
                 return None, None  # Signal auth failure
             return None, None
 
         logger.info(
-            f"WebSocket token found in cookie for app '{app_slug}' " "(cookie-based authentication)"
+            f"WebSocket token found in cookie for app '{app_slug}' "
+            "(cookie-based authentication, fallback)"
         )
 
         # Decode and validate token
@@ -428,8 +476,11 @@ async def authenticate_websocket(
                 f"(method: cookie)"
             )
             return user_id, user_email
-        except (jwt.ExpiredSignatureError, jwt.InvalidTokenError) as decode_error:
-            logger.error(f"JWT decode error for app '{app_slug}': {decode_error}", exc_info=True)
+        except (jwt.ExpiredSignatureError, jwt.InvalidTokenError):
+            logger.exception(
+                f"❌ JWT decode error for app '{app_slug}'. "
+                f"Token present: {bool(token)}, Token length: {len(token) if token else 0}"
+            )
             raise
 
     except WebSocketDisconnect:
@@ -689,12 +740,26 @@ def create_websocket_endpoint(
             # CRITICAL: Authenticate BEFORE accepting connection
             # This prevents CSRF middleware from rejecting established connections
             # We can access headers/query_params before accept() is called
+
+            # Debug: Log cookies before authentication
+            try:
+                cookies = _get_cookies_from_websocket(websocket)
+                cookie_names = list(cookies.keys()) if cookies else []
+                logger.info(
+                    f"🔍 WebSocket cookies for app '{app_slug}': {cookie_names} "
+                    f"(require_auth={require_auth})"
+                )
+            except (AttributeError, TypeError, KeyError, RuntimeError) as cookie_error:
+                logger.warning(f"Could not extract cookies for debugging: {cookie_error}")
+
             user_id, user_email = await authenticate_websocket(websocket, app_slug, require_auth)
 
             # Handle authentication failure
             if require_auth and not user_id:
-                logger.warning(
-                    f"WebSocket authentication failed for app '{app_slug}' - rejecting connection"
+                logger.error(
+                    f"❌ WebSocket authentication FAILED for app '{app_slug}' - "
+                    f"rejecting connection. require_auth={require_auth}, "
+                    f"user_id={user_id}, user_email={user_email}"
                 )
                 # Reject without accepting - FastAPI will send 403 if accept() not called
                 # We can't call websocket.close() before accept(), so we just return