mdb-engine 0.5.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

mdb_engine/auth/provider.py

@@ -86,6 +86,16 @@ class CasbinAdapter(BaseAuthorizationProvider):
         self._cache_lock = asyncio.Lock()
         self._mark_initialized()
 
+    @property
+    def enforcer(self):
+        """
+        Get the Casbin enforcer instance.
+
+        Returns:
+            casbin.AsyncEnforcer instance
+        """
+        return self._enforcer
+
     async def check(
         self,
         subject: str,

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:
@@ -652,6 +682,47 @@ class SharedUserPool:
         )
         return result.modified_count > 0
 
+    async def update_user_metadata(
+        self,
+        email: str,
+        metadata: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """
+        Update user metadata fields.
+
+        This allows adding or updating custom fields on the user document
+        beyond the core schema (e.g., name, profile data, preferences).
+
+        Args:
+            email: User email
+            metadata: Dictionary of fields to update
+                (e.g., {"name": "John Doe", "preferences": {...}})
+
+        Returns:
+            Updated user document (without password_hash) or None if user not found
+
+        Example:
+            user = await pool.update_user_metadata(
+                "user@example.com",
+                {"name": "John Doe", "phone": "+1234567890"}
+            )
+        """
+        # Build update document, ensuring updated_at is always set
+        update_doc = {"$set": {**metadata, "updated_at": datetime.utcnow()}}
+
+        result = await self._collection.update_one(
+            {"email": email},
+            update_doc,
+        )
+
+        if result.modified_count > 0:
+            # Fetch and return updated user
+            updated_user = await self._collection.find_one({"email": email})
+            if updated_user:
+                logger.info(f"Updated metadata for user '{email}': {list(metadata.keys())}")
+                return self._sanitize_user(updated_user)
+        return None
+
     @staticmethod
     def user_has_role(
         user: dict[str, Any],

mdb_engine/auth/users.py

@@ -1376,7 +1376,8 @@ async def sync_app_user_to_casbin(
             logger.debug("sync_app_user_to_casbin: Provider is not CasbinAdapter, skipping")
             return False
 
-        enforcer = authz_provider._enforcer
+        # Access enforcer via property if available, otherwise fallback to private member
+        enforcer = getattr(authz_provider, "enforcer", None) or authz_provider._enforcer  # noqa: SLF001
 
         # Get user ID
         user_id = str(user.get("_id") or user.get("app_user_id", ""))

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(

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