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

mdb_engine/__init__.py

@@ -82,15 +82,19 @@ from .repositories import Entity, MongoRepository, Repository, UnitOfWork
 from .utils import clean_mongo_doc, clean_mongo_docs
 
 __version__ = (
-    "0.5.0"  # Major WebSocket security overhaul - cookie-based authentication
-    # - BREAKING CHANGE: Removed subprotocol tunneling support
-    # - NEW: Exclusive httpOnly cookie-based WebSocket authentication
-    # - Comprehensive security guide (WEBSOCKET_SECURITY_MULTI_APP_SSO.md)
-    # - Updated all documentation to reflect cookie-based authentication
-    # - Enhanced CSRF protection for WebSocket upgrades
-    # - Added integration tests for cookie-based WebSocket authentication
-    # - Complete test coverage for WebSocket security scenarios
-    # - Multi-app SSO compatibility with path="/" cookies
+    "0.6.0"  # Secure-by-default WebSocket authentication with encrypted session keys
+    # - NEW: WebSocket session key generation and management
+    # - NEW: Envelope encryption for WebSocket session keys
+    # - NEW: Secure-by-default CSRF protection (csrf_required: true)
+    # - NEW: WebSocketSessionManager with private collection storage
+    # - NEW: Session key endpoint (/auth/websocket-session)
+    # - NEW: Session key integration in login flow
+    # - ENHANCED: WebSocket authentication with session key support
+    # - ENHANCED: CSRF middleware session key validation
+    # - ENHANCED: Multi-app WebSocket routing with session keys
+    # - BACKWARD COMPATIBLE: Cookie-based authentication fallback
+    # - UPDATED: All documentation for secure-by-default approach
+    # - COMPREHENSIVE: Unit and integration tests for session keys
 )
 
 __all__ = [

mdb_engine/auth/__init__.py

@@ -125,6 +125,12 @@ from .utils import (
     validate_password_strength_async,
 )
 
+# WebSocket sessions
+from .websocket_sessions import (
+    WebSocketSessionManager,
+    create_websocket_session_endpoint,
+)
+
 __all__ = [
     # Base classes
     "BaseAuthorizationProvider",
@@ -232,4 +238,7 @@ __all__ = [
     "generate_csrf_token",
     "validate_csrf_token",
     "get_csrf_token",
+    # WebSocket sessions
+    "WebSocketSessionManager",
+    "create_websocket_session_endpoint",
 ]

mdb_engine/auth/csrf.py

@@ -102,14 +102,16 @@ def validate_csrf_token(
         try:
             parts = token.split(":")
             if len(parts) != 3:
+                logger.debug("CSRF token has wrong format (expected 3 parts)")
                 return False
 
             raw_token, timestamp_str, signature = parts
             timestamp = int(timestamp_str)
 
             # Check age
-            if time.time() - timestamp > max_age:
-                logger.debug("CSRF token expired")
+            age = time.time() - timestamp
+            if age > max_age:
+                logger.debug(f"CSRF token expired (age: {age:.0f}s, max: {max_age}s)")
                 return False
 
             # Verify signature
@@ -119,7 +121,10 @@ def validate_csrf_token(
             ]
 
             if not hmac.compare_digest(signature, expected_sig):
-                logger.warning("CSRF token signature mismatch")
+                logger.warning(
+                    f"CSRF token signature mismatch. "
+                    f"Token format: signed, Has secret: {bool(secret)}"
+                )
                 return False
 
             return True
@@ -128,7 +133,10 @@ def validate_csrf_token(
             return False
 
     # Simple token validation (just check it exists and has reasonable length)
-    return len(token) >= CSRF_TOKEN_LENGTH
+    is_valid = len(token) >= CSRF_TOKEN_LENGTH
+    if not is_valid:
+        logger.debug(f"CSRF token too short (length: {len(token)}, required: {CSRF_TOKEN_LENGTH})")
+    return is_valid
 
 
 class CSRFMiddleware(BaseHTTPMiddleware):
@@ -197,10 +205,116 @@ class CSRFMiddleware(BaseHTTPMiddleware):
                 return True
         return False
 
+    def _websocket_requires_csrf(self, request: Request, path: str) -> bool:
+        """
+        Check if WebSocket endpoint requires CSRF validation.
+
+        Defaults to True (security by default). Can be disabled per-endpoint via manifest.json:
+        websockets.{endpoint}.auth.csrf_required = false
+
+        Args:
+            request: FastAPI request
+            path: WebSocket path (e.g., "/app-3/ws")
+
+        Returns:
+            True if CSRF validation is required, False otherwise
+        """
+        # Check parent app state for WebSocket configs
+        websocket_configs = getattr(request.app.state, "websocket_configs", None)
+        if not websocket_configs:
+            # No WebSocket configs found - use default (CSRF required for security by default)
+            return True
+
+        # Normalize path for matching
+        normalized_path = path.rstrip("/")
+
+        # Try to find matching app config
+        # WebSocket paths are registered as /app-slug/endpoint-path
+        # e.g., /app-3/ws where app_slug="app-3" and endpoint_path="/ws"
+        for app_slug, config in websocket_configs.items():
+            # Check each endpoint in this app's config
+            for endpoint_name, endpoint_config in config.items():
+                endpoint_path = endpoint_config.get("path", "")
+                # Normalize endpoint path
+                normalized_endpoint = endpoint_path.rstrip("/")
+
+                # Match patterns:
+                # 1. Full path match: /app-slug/endpoint-path
+                # 2. Endpoint-only match: /endpoint-path (if path starts with endpoint)
+                expected_full_path = f"/{app_slug}{normalized_endpoint}"
+                if (
+                    normalized_path == expected_full_path
+                    or normalized_path.endswith(normalized_endpoint)
+                    or normalized_path == normalized_endpoint
+                ):
+                    auth_config = endpoint_config.get("auth", {})
+                    if isinstance(auth_config, dict):
+                        # Return csrf_required setting (defaults to True - security by default)
+                        csrf_required = auth_config.get("csrf_required", True)
+                        logger.debug(
+                            f"WebSocket {path} csrf_required={csrf_required} "
+                            f"(from app={app_slug}, endpoint={endpoint_name})"
+                        )
+                        return csrf_required
+
+        # No matching config found - use default (CSRF required for security by default)
+        logger.debug(f"No WebSocket config match for {path}, using default csrf_required=true")
+        return True
+
     def _is_websocket_upgrade(self, request: Request) -> bool:
         """Check if request is a WebSocket upgrade request."""
         upgrade_header = request.headers.get("upgrade", "").lower()
-        return upgrade_header == "websocket"
+        connection_header = request.headers.get("connection", "").lower()
+
+        # Primary check: WebSocket upgrade requires both Upgrade: websocket
+        # and Connection: Upgrade headers
+        has_upgrade_header = upgrade_header == "websocket"
+        has_connection_upgrade = "upgrade" in connection_header or "websocket" in connection_header
+
+        # Secondary check: If upgrade header is present but connection is
+        # overridden (e.g., by TestClient), check if path matches a known
+        # WebSocket route pattern
+        path_matches_websocket_route = False
+        if has_upgrade_header and not has_connection_upgrade:
+            # Check if path matches any configured WebSocket route
+            websocket_configs = getattr(request.app.state, "websocket_configs", None)
+            if websocket_configs:
+                path = request.url.path.rstrip("/") or "/"
+                for app_slug, config in websocket_configs.items():
+                    for _endpoint_name, endpoint_config in config.items():
+                        endpoint_path = endpoint_config.get("path", "").rstrip("/") or "/"
+                        # Try various path matching patterns
+                        expected_full_path = (
+                            f"/{app_slug}{endpoint_path}"
+                            if endpoint_path != "/"
+                            else f"/{app_slug}"
+                        )
+                        # Match patterns:
+                        # 1. Exact match with app prefix: /app-slug/endpoint-path
+                        # 2. Endpoint-only match: /endpoint-path (if path ends with endpoint)
+                        # 3. Root match: / matches / or /app-slug
+                        if (
+                            path == expected_full_path
+                            or path.endswith(endpoint_path)
+                            or path == endpoint_path
+                            or (path == "/" and endpoint_path == "/")
+                            or (path == f"/{app_slug}" and endpoint_path == "/")
+                        ):
+                            path_matches_websocket_route = True
+                            break
+                    if path_matches_websocket_route:
+                        break
+
+        is_websocket = has_upgrade_header and (
+            has_connection_upgrade or path_matches_websocket_route
+        )
+        if is_websocket:
+            logger.debug(
+                f"WebSocket upgrade detected: path={request.url.path}, "
+                f"upgrade={upgrade_header}, connection={connection_header}, "
+                f"path_match={path_matches_websocket_route}"
+            )
+        return is_websocket
 
     def _get_allowed_origins(self, request: Request) -> list[str]:
         """
@@ -298,9 +412,22 @@ class CSRFMiddleware(BaseHTTPMiddleware):
         path = request.url.path
         method = request.method
 
+        # Debug: Log all requests to see what's happening
+        upgrade_header = request.headers.get("upgrade", "").lower()
+        connection_header = request.headers.get("connection", "").lower()
+        if upgrade_header or "websocket" in path.lower():
+            logger.info(
+                f"🔍 CSRF middleware: {method} {path}, "
+                f"upgrade={upgrade_header}, connection={connection_header}"
+            )
+
         # CRITICAL: Handle WebSocket upgrade requests BEFORE other CSRF checks
         # WebSocket upgrades use cookie-based authentication and require CSRF validation
         if self._is_websocket_upgrade(request):
+            logger.info(
+                f"🔌 CSRF middleware processing WebSocket upgrade: {path}, "
+                f"origin: {request.headers.get('origin')}"
+            )
             # Always validate origin for WebSocket connections (CSWSH protection)
             if not self._validate_websocket_origin(request):
                 logger.warning(
@@ -315,51 +442,160 @@ class CSRFMiddleware(BaseHTTPMiddleware):
 
             # Cookie-based authentication requires CSRF protection
             # Check if authentication token cookie is present
-            auth_token_cookie = request.cookies.get("token")
+            # Use same cookie name as SharedAuthMiddleware for consistency
+            from .shared_middleware import AUTH_COOKIE_NAME
+
+            auth_token_cookie = request.cookies.get(AUTH_COOKIE_NAME)
             if auth_token_cookie:
-                # For WebSocket upgrades, CSRF protection relies on:
+                # SECURITY BY DEFAULT: WebSocket CSRF protection uses encrypted session keys
+                # stored in private collection via envelope encryption.
+                #
+                # Security Model:
                 # 1. Origin validation (already done above) - primary defense
-                # 2. SameSite cookies - prevents cross-site cookie sending
-                # 3. CSRF cookie presence - ensures session is established
+                # 2. Encrypted session key validation - CSRF protection via database
+                # 3. SameSite cookies - prevents cross-site cookie sending
                 #
-                # Note: JavaScript WebSocket API cannot set custom headers,
-                # so we cannot use double-submit cookie pattern (cookie + header).
-                # Instead, we rely on Origin validation + SameSite cookies for CSRF protection.
-                csrf_cookie_token = request.cookies.get(self.cookie_name)
-                if not csrf_cookie_token:
-                    logger.warning(f"WebSocket upgrade missing CSRF cookie for {path}")
-                    return JSONResponse(
-                        status_code=status.HTTP_403_FORBIDDEN,
-                        content={"detail": "CSRF token missing for WebSocket authentication"},
+                # Session keys are:
+                # - Generated on authentication
+                # - Encrypted using envelope encryption (same as app secrets)
+                # - Stored in _mdb_engine_websocket_sessions private collection
+                # - Validated during WebSocket upgrade
+
+                # Check if this WebSocket endpoint requires CSRF validation
+                csrf_required = self._websocket_requires_csrf(request, path)
+
+                if csrf_required:
+                    # Try to get WebSocket session manager from app state
+                    websocket_session_manager = getattr(
+                        request.app.state, "websocket_session_manager", None
                     )
 
-                # Validate CSRF token signature if secret is used
-                if self.secret and not validate_csrf_token(
-                    csrf_cookie_token, self.secret, self.token_ttl
-                ):
-                    logger.warning(f"WebSocket CSRF token validation failed for {path}")
-                    return JSONResponse(
-                        status_code=status.HTTP_403_FORBIDDEN,
-                        content={
-                            "detail": "CSRF token expired or invalid for WebSocket connection"
-                        },
+                    if websocket_session_manager:
+                        # Use encrypted session key validation (secure-by-default)
+                        session_key = request.query_params.get(
+                            "session_key"
+                        ) or request.headers.get("X-WebSocket-Session-Key")
+
+                        if not session_key:
+                            logger.error(
+                                f"❌ WebSocket upgrade missing session key for {path}. "
+                                f"Auth cookie present: {bool(auth_token_cookie)}. "
+                                f"Tip: Generate session key via /auth/websocket-session endpoint."
+                            )
+                            return JSONResponse(
+                                status_code=status.HTTP_403_FORBIDDEN,
+                                content={
+                                    "detail": (
+                                        "WebSocket session key missing. "
+                                        "Generate session key via /auth/websocket-session endpoint."
+                                    )
+                                },
+                            )
+
+                        # Validate session key against encrypted storage
+                        try:
+                            session_data = await websocket_session_manager.validate_session(
+                                session_key
+                            )
+                            if not session_data:
+                                logger.error(
+                                    f"❌ WebSocket session key validation failed for {path}. "
+                                    f"Session key: {session_key[:16]}..."
+                                )
+                                return JSONResponse(
+                                    status_code=status.HTTP_403_FORBIDDEN,
+                                    content={
+                                        "detail": (
+                                            "WebSocket session key expired or invalid. "
+                                            "Generate a new session key."
+                                        )
+                                    },
+                                )
+
+                            # Store session data in request state for WebSocket handler
+                            request.state.websocket_session = session_data
+                            logger.debug(
+                                f"✅ WebSocket session key validated for {path} "
+                                f"(user: {session_data.get('user_id')})"
+                            )
+                        except (
+                            ValueError,
+                            TypeError,
+                            AttributeError,
+                            RuntimeError,
+                        ):
+                            logger.exception("Error validating WebSocket session key")
+                            return JSONResponse(
+                                status_code=status.HTTP_403_FORBIDDEN,
+                                content={"detail": "WebSocket session validation error"},
+                            )
+                    else:
+                        # Fallback to cookie-based CSRF (backward compatibility)
+                        csrf_cookie_token = request.cookies.get(self.cookie_name)
+                        if not csrf_cookie_token:
+                            logger.error(
+                                f"❌ WebSocket upgrade missing CSRF cookie for {path}. "
+                                f"Auth cookie present: {bool(auth_token_cookie)}, "
+                                f"CSRF cookie name: {self.cookie_name}, "
+                                f"Available cookies: {list(request.cookies.keys())}. "
+                                f"Tip: Make a GET request first to receive CSRF cookie."
+                            )
+                            return JSONResponse(
+                                status_code=status.HTTP_403_FORBIDDEN,
+                                content={
+                                    "detail": (
+                                        "CSRF token missing for WebSocket authentication. "
+                                        "Make a GET request first to receive the CSRF cookie."
+                                    )
+                                },
+                            )
+
+                        # Validate CSRF token signature if secret is used
+                        if self.secret and not validate_csrf_token(
+                            csrf_cookie_token, self.secret, self.token_ttl
+                        ):
+                            logger.error(f"❌ WebSocket CSRF token validation failed for {path}.")
+                            return JSONResponse(
+                                status_code=status.HTTP_403_FORBIDDEN,
+                                content={
+                                    "detail": (
+                                        "CSRF token expired or invalid for WebSocket connection"
+                                    )
+                                },
+                            )
+
+                        # If CSRF header is provided, validate it matches the cookie
+                        # (Header is optional for WebSocket, but if present, must match cookie)
+                        csrf_header_token = request.headers.get(self.header_name)
+                        if csrf_header_token:
+                            if not hmac.compare_digest(csrf_cookie_token, csrf_header_token):
+                                logger.error(
+                                    f"❌ WebSocket CSRF header mismatch for {path}. "
+                                    f"Cookie token and header token do not match."
+                                )
+                                return JSONResponse(
+                                    status_code=status.HTTP_403_FORBIDDEN,
+                                    content={
+                                        "detail": (
+                                            "CSRF token mismatch: header token does not "
+                                            "match cookie token"
+                                        )
+                                    },
+                                )
+                            logger.debug(
+                                f"✅ CSRF header validated and matches cookie for WebSocket {path}"
+                            )
+
+                        logger.debug(f"✅ CSRF cookie validation passed for WebSocket {path}")
+                else:
+                    logger.debug(
+                        f"✅ WebSocket CSRF validation skipped for {path} "
+                        f"(csrf_required=false, Origin validation sufficient)"
                     )
 
-                # Optional: If CSRF header is provided, validate it matches cookie
-                # (Some clients may send it, but it's not required for WebSocket upgrades)
-                header_token = request.headers.get(self.header_name)
-                if header_token:
-                    # If header is provided, validate it matches cookie (double-submit pattern)
-                    if not hmac.compare_digest(csrf_cookie_token, header_token):
-                        logger.warning(f"WebSocket CSRF token mismatch for {path}")
-                        return JSONResponse(
-                            status_code=status.HTTP_403_FORBIDDEN,
-                            content={"detail": "CSRF token invalid for WebSocket connection"},
-                        )
-
                 logger.debug(
                     f"WebSocket upgrade CSRF validation passed for {path} "
-                    f"(Origin validated, CSRF cookie present)"
+                    f"(Origin validated, CSRF validated)"
                 )
 
             # Origin validated (and CSRF validated if authenticated)

mdb_engine/auth/shared_users.py

@@ -120,6 +120,7 @@ class SharedUserPool:
         token_expiry_hours: int = DEFAULT_TOKEN_EXPIRY_HOURS,
         allow_insecure_dev: bool = False,
         blacklist_fail_closed: bool = True,
+        websocket_session_manager: Any | None = None,
     ):
         """
         Initialize the shared user pool.
@@ -174,6 +175,7 @@ class SharedUserPool:
 
         self._token_expiry_hours = token_expiry_hours
         self._blacklist_indexes_created = False
+        self._websocket_session_manager = websocket_session_manager
 
         logger.info(f"SharedUserPool initialized (algorithm={jwt_algorithm})")
 
@@ -340,7 +342,9 @@ class SharedUserPool:
         ip_address: str | None = None,
         fingerprint: str | None = None,
         session_binding: dict[str, Any] | None = None,
-    ) -> str | None:
+        generate_websocket_session: bool = True,
+        app_slug: str | None = None,
+    ) -> str | tuple[str, str] | None:
         """
         Authenticate user and return JWT token.
 
@@ -352,9 +356,14 @@ class SharedUserPool:
             session_binding: Session binding config from manifest:
                 - bind_ip: Include IP in token claims
                 - bind_fingerprint: Include fingerprint in token claims
+            generate_websocket_session: If True and WebSocket session manager available,
+                                       also generate WebSocket session key (default: True)
+            app_slug: Optional app slug for WebSocket session scoping
 
         Returns:
-            JWT token if authentication succeeds, None otherwise
+            JWT token if authentication succeeds, None otherwise.
+            If generate_websocket_session=True and session manager available,
+            returns tuple (jwt_token, websocket_session_key), otherwise just jwt_token.
         """
         user = await self._collection.find_one(
             {
@@ -392,7 +401,28 @@ class SharedUserPool:
         # Generate JWT token with session binding claims
         token = self._generate_token(user, extra_claims=extra_claims or None)
 
+        # Generate WebSocket session key if requested and manager available
+        websocket_session_key = None
+        if generate_websocket_session and self._websocket_session_manager:
+            try:
+                user_id = str(user["_id"])
+                websocket_session_key = await self._websocket_session_manager.create_session(
+                    user_id=user_id,
+                    user_email=email,
+                    app_slug=app_slug,
+                )
+                logger.debug(
+                    f"Generated WebSocket session key for user '{email}' " f"(app: {app_slug})"
+                )
+            except (ValueError, TypeError, AttributeError, RuntimeError) as e:
+                # Log but don't fail authentication if WebSocket session generation fails
+                logger.warning(f"Failed to generate WebSocket session key: {e}")
+
         logger.info(f"User '{email}' authenticated successfully")
+
+        # Return tuple if WebSocket session key was generated, otherwise just token
+        if websocket_session_key:
+            return (token, websocket_session_key)
         return token
 
     async def validate_token(self, token: str) -> dict[str, Any] | None:

mdb_engine/auth/utils.py

@@ -514,16 +514,41 @@ async def login_user(
                 ip_address=device_info.get("ip_address"),
             )
 
+        # Generate WebSocket session key if WebSocket session manager available
+        websocket_session_key = None
+        try:
+            # Try to get WebSocket session manager from app state
+            app = getattr(request, "app", None)
+            if app:
+                websocket_session_manager = getattr(app.state, "websocket_session_manager", None)
+                if websocket_session_manager:
+                    # Get app slug from request state if available
+                    app_slug = getattr(request.state, "app_slug", None)
+                    websocket_session_key = await websocket_session_manager.create_session(
+                        user_id=str(user["_id"]),
+                        user_email=user["email"],
+                        app_slug=app_slug,
+                    )
+                    logger.debug(
+                        f"Generated WebSocket session key for user '{user['email']}' "
+                        f"(app: {app_slug})"
+                    )
+        except (ValueError, TypeError, AttributeError, RuntimeError) as e:
+            # Log but don't fail login if WebSocket session generation fails
+            logger.warning(f"Failed to generate WebSocket session key during login: {e}")
+
         # Create response
+        response_data = {
+            "success": True,
+            "user": {"email": user["email"], "user_id": str(user["_id"])},
+        }
+        if websocket_session_key:
+            response_data["websocket_session_key"] = websocket_session_key
+
         if redirect_url:
             response = RedirectResponse(url=redirect_url, status_code=302)
         else:
-            response = JSONResponse(
-                {
-                    "success": True,
-                    "user": {"email": user["email"], "user_id": str(user["_id"])},
-                }
-            )
+            response = JSONResponse(response_data)
 
         # Set cookies
         set_auth_cookies(