ragbits-chat 1.4.0.dev202512160238__py3-none-any.whl → 1.4.0.dev202601010248__py3-none-any.whl

ragbits/chat/auth/backends.py

@@ -1,14 +1,21 @@
+import logging
+import secrets
 import uuid
 from datetime import datetime, timedelta, timezone
 from typing import Any, cast
+from urllib.parse import urlencode
 
 import bcrypt
-from jose import jwt
-from jose.exceptions import ExpiredSignatureError, JWTError
+import httpx
 
 from ragbits.chat.auth.base import AuthenticationBackend, AuthenticationResponse, AuthOptions
-from ragbits.chat.auth.types import JWTToken, OAuth2Credentials, User, UserCredentials
-from ragbits.core.utils import get_secret_key
+from ragbits.chat.auth.oauth2_providers import OAuth2Provider
+from ragbits.chat.auth.types import OAuth2Credentials, Session, SessionStore, User, UserCredentials
+
+logger = logging.getLogger(__name__)
+
+# Minimum length for session ID truncation in logs (for readability while maintaining some privacy)
+SESSION_ID_LOG_LENGTH = 8
 
 
 class ListAuthenticationBackend(AuthenticationBackend):
@@ -17,6 +24,8 @@ class ListAuthenticationBackend(AuthenticationBackend):
     def __init__(
         self,
         users: list[dict[str, Any]],
+        session_store: SessionStore,
+        session_expiry_hours: int = 24,
         default_options: AuthOptions | None = None,
     ):
         """
@@ -24,17 +33,16 @@ class ListAuthenticationBackend(AuthenticationBackend):
 
         Args:
             users: List of user dicts with 'username', 'password', and optional fields
-            jwt_secret: Secret key for JWT jwt_token signing (generates random if not provided)
+            session_store: Session storage backend
+            session_expiry_hours: Hours until session expires (default: 24)
             default_options: Default options for the component
         """
         if default_options is None:
             default_options = AuthOptions()
         super().__init__(default_options)
         self.users = {}
-        self.jwt_secret = get_secret_key()
-        self.jwt_algorithm = default_options.jwt_algorithm
-        self.token_expiry_minutes = default_options.token_expiry_minutes
-        self.revoked_tokens: set[str] = set()  # Blacklist for revoked tokens
+        self.session_store = session_store
+        self.session_expiry_hours = session_expiry_hours
 
         for user_data in users:
             # Hash passwords with bcrypt for security
@@ -51,7 +59,7 @@ class ListAuthenticationBackend(AuthenticationBackend):
                 ),
             }
 
-    async def authenticate_with_credentials(self, credentials: UserCredentials) -> AuthenticationResponse:  # noqa: PLR6301
+    async def authenticate_with_credentials(self, credentials: UserCredentials) -> AuthenticationResponse:
         """
         Authenticate into backend using provided credentials
 
@@ -60,118 +68,548 @@ class ListAuthenticationBackend(AuthenticationBackend):
         Returns:
             AuthenticationResponse: Result of authentication
         """
+        logger.debug("Attempting credential authentication for user: %s", credentials.username)
+
         user_data = self.users.get(credentials.username)
         if not user_data:
+            logger.warning("Authentication failed: user '%s' not found", credentials.username)
             return AuthenticationResponse(success=False, error_message="User not found")
 
         # Verify password with bcrypt
         password_hash = str(user_data["password_hash"])
         if not bcrypt.checkpw(credentials.password.encode("utf-8"), password_hash.encode("utf-8")):
+            logger.warning("Authentication failed: invalid password for user '%s'", credentials.username)
             return AuthenticationResponse(success=False, error_message="Invalid password")
 
         user = cast(User, user_data["user"])
 
-        # Create JWT jwt_token
-        jwt_token = self._create_jwt_token(user)
+        # Create session
+        now = datetime.now(timezone.utc)
+        session = Session(
+            session_id="",  # Will be generated by session store
+            user=user,
+            provider="credentials",
+            oauth_token="",  # Not applicable for credentials
+            token_type="",
+            created_at=now,
+            expires_at=now + timedelta(hours=self.session_expiry_hours),
+        )
+        session_id = await self.session_store.create_session(session)
+        logger.info("User '%s' authenticated successfully, session created", credentials.username)
+        return AuthenticationResponse(success=True, user=user, session_id=session_id)
+
+    async def validate_session(self, session_id: str) -> AuthenticationResponse:
+        """
+        Validate a session.
 
-        return AuthenticationResponse(success=True, user=user, jwt_token=jwt_token)
+        Args:
+            session_id: The session ID to validate
 
-    def _create_jwt_token(self, user: User) -> JWTToken:
-        """Create a JWT jwt_token for the user."""
-        now = datetime.now(timezone.utc)
-        expires_in = self.token_expiry_minutes * 60  # Convert to seconds
-
-        payload = {
-            "user_id": user.user_id,
-            "username": user.username,
-            "email": user.email,
-            "full_name": user.full_name,
-            "roles": user.roles,
-            "metadata": user.metadata,
-            "iat": now,
-            "exp": now + timedelta(minutes=self.token_expiry_minutes),
+        Returns:
+            AuthenticationResponse with user if valid
+        """
+        logger.debug(
+            "Validating session: %s...",
+            session_id[:SESSION_ID_LOG_LENGTH] if len(session_id) >= SESSION_ID_LOG_LENGTH else session_id,
+        )
+        session = await self.session_store.get_session(session_id)
+
+        if not session:
+            logger.debug("Session validation failed: session not found or expired")
+            return AuthenticationResponse(success=False, error_message="Invalid or expired session")
+
+        logger.debug("Session validated successfully for user: %s", session.user.username)
+        return AuthenticationResponse(success=True, user=session.user)
+
+    async def authenticate_with_oauth2(  # noqa: PLR6301
+        self, oauth_credentials: OAuth2Credentials
+    ) -> AuthenticationResponse:
+        """
+        Authenticate user with OAuth2 credentials.
+
+        Args:
+            oauth_credentials: OAuth2 credentials
+
+        Returns:
+            AuthenticationResponse: Authentication failure as OAuth2 is not supported
+        """
+        return AuthenticationResponse(success=False, error_message="OAuth2 not supported by ListAuthentication")
+
+    async def revoke_session(self, session_id: str) -> bool:
+        """
+        Revoke a session.
+
+        Args:
+            session_id: The session ID to revoke
+
+        Returns:
+            True if session was revoked
+        """
+        logger.debug(
+            "Revoking session: %s...",
+            session_id[:SESSION_ID_LOG_LENGTH] if len(session_id) >= SESSION_ID_LOG_LENGTH else session_id,
+        )
+        success = await self.session_store.delete_session(session_id)
+        if success:
+            logger.info("Session revoked successfully")
+        else:
+            logger.debug("Session revocation failed: session not found")
+        return success
+
+
+class OAuth2AuthenticationBackend(AuthenticationBackend):
+    """Generic OAuth2 authentication backend supporting multiple providers."""
+
+    def __init__(
+        self,
+        session_store: SessionStore,
+        provider: OAuth2Provider,
+        client_id: str | None = None,
+        client_secret: str | None = None,
+        redirect_uri: str | None = None,
+        session_expiry_hours: int = 24,
+        default_options: AuthOptions | None = None,
+    ):
+        """
+        Initialize OAuth2 authentication backend.
+
+        Args:
+            session_store: Session storage backend
+            provider: OAuth2 provider implementation (e.g., DiscordOAuth2Provider, GoogleOAuth2Provider)
+            client_id: OAuth2 client ID (or set {PROVIDER}_CLIENT_ID env var)
+            client_secret: OAuth2 client secret (or set {PROVIDER}_CLIENT_SECRET env var)
+            redirect_uri: Callback URL for OAuth2 flow (or set OAUTH2_REDIRECT_URI env var,
+                         defaults to 'http://localhost:8000/api/auth/callback/{provider_name}')
+            session_expiry_hours: Hours until session expires (default: 24)
+            default_options: Default options for the component
+
+        Note:
+            The default redirect_uri uses a provider-specific path for better isolation and debugging.
+            For Discord, it defaults to: http://localhost:8000/api/auth/callback/discord
+        """
+        import os
+
+        if default_options is None:
+            default_options = AuthOptions()
+        super().__init__(default_options)
+
+        self.session_store = session_store
+        self.session_expiry_hours = session_expiry_hours
+        self.provider = provider
+
+        # Get credentials from args or environment variables
+        self.client_id = client_id or os.getenv(f"{self.provider.name.upper()}_CLIENT_ID")
+        self.client_secret = client_secret or os.getenv(f"{self.provider.name.upper()}_CLIENT_SECRET")
+
+        # Use provider-specific callback URL for better isolation and debugging
+        if not redirect_uri:
+            # Get base URL from environment variable or use default
+            base_url = os.getenv("OAUTH2_CALLBACK_BASE_URL", "http://localhost:8000")
+            # Remove trailing slash from base URL
+            base_url = base_url.rstrip("/")
+            # Construct redirect URI with base URL and provider name
+            redirect_uri = f"{base_url}/api/auth/callback/{self.provider.name}"
+
+        # remove trailing slash from redirect URI
+        redirect_uri = redirect_uri.rstrip("/")
+        # set redirect URI on the backend
+        self.redirect_uri = redirect_uri
+
+        if not self.client_id or not self.client_secret:
+            raise ValueError(
+                f"OAuth2 credentials not provided. Either pass client_id and client_secret to the constructor, "
+                f"or set {self.provider.name.upper()}_CLIENT_ID and {self.provider.name.upper()}_CLIENT_SECRET "
+                f"environment variables."
+            )
+
+        # State storage for CSRF protection (in production, use Redis or similar)
+        self.pending_states: dict[str, datetime] = {}
+
+    def generate_authorize_url(self) -> tuple[str, str]:
+        """
+        Generate OAuth2 authorization URL with state parameter.
+
+        Returns:
+            Tuple of (authorize_url, state)
+        """
+        state = secrets.token_urlsafe(32)
+        self.pending_states[state] = datetime.now(timezone.utc)
+
+        params = {
+            "client_id": self.client_id,
+            "redirect_uri": self.redirect_uri,
+            "response_type": "code",
+            "scope": self.provider.scope,
+            "state": state,
         }
 
-        access_token = jwt.encode(payload, self.jwt_secret, algorithm=self.jwt_algorithm)
+        authorize_url = f"{self.provider.authorize_url}?{urlencode(params)}"
+        logger.debug("Generated OAuth2 authorization URL for provider '%s'", self.provider.name)
+        return authorize_url, state
+
+    def verify_state(self, state: str) -> bool:
+        """
+        Verify OAuth2 state parameter for CSRF protection.
+
+        Args:
+            state: State parameter to verify
+
+        Returns:
+            True if state is valid, False otherwise
+        """
+        if state not in self.pending_states:
+            logger.warning("OAuth2 state verification failed: state not found (possible CSRF attempt)")
+            return False
+
+        # Check if state is not expired (valid for 10 minutes)
+        created_at = self.pending_states[state]
+        if datetime.now(timezone.utc) - created_at > timedelta(minutes=10):
+            del self.pending_states[state]
+            logger.warning("OAuth2 state verification failed: state expired")
+            return False
+
+        # Remove state after verification (one-time use)
+        del self.pending_states[state]
+        logger.debug("OAuth2 state verified successfully")
+        return True
+
+    def cleanup_expired_states(self) -> int:
+        """
+        Remove expired state tokens from storage.
+
+        This method removes state tokens that are older than 10 minutes to prevent
+        memory leaks from abandoned OAuth2 flows.
+
+        Returns:
+            Number of state tokens removed
+
+        Example:
+            To schedule periodic cleanup:
+
+            ```python
+            import asyncio
+
+
+            async def cleanup_loop():
+                while True:
+                    await asyncio.sleep(600)  # Run every 10 minutes
+                    removed = oauth2_backend.cleanup_expired_states()
+                    if removed > 0:
+                        logger.info(f"Cleaned up {removed} expired OAuth2 states")
+            ```
+        """
+        now = datetime.now(timezone.utc)
+        states_to_remove = []
+
+        # Find expired states (older than 10 minutes)
+        for state, created_at in list(self.pending_states.items()):
+            if now - created_at > timedelta(minutes=10):
+                states_to_remove.append(state)
 
-        token_type = "bearer"  # noqa: S105
-        return JWTToken(access_token=access_token, token_type=token_type, expires_in=expires_in, user=user)
+        # Remove expired states
+        for state in states_to_remove:
+            self.pending_states.pop(state, None)
 
-    async def validate_token(self, token: str) -> AuthenticationResponse:
-        """Validate a JWT jwt_token."""
-        # Check if token is blacklisted (revoked)
-        if token in self.revoked_tokens:
-            return AuthenticationResponse(success=False, error_message="Token has been revoked")
+        if states_to_remove:
+            logger.info("Cleaned up %d expired OAuth2 state tokens", len(states_to_remove))
 
+        return len(states_to_remove)
+
+    async def exchange_code_for_token(self, code: str) -> str | None:
+        """
+        Exchange authorization code for access token.
+
+        Args:
+            code: Authorization code from OAuth2 provider
+
+        Returns:
+            Access token if successful, None otherwise
+        """
+        logger.debug("Exchanging authorization code for token with provider '%s'", self.provider.name)
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    self.provider.token_url,
+                    data={
+                        "client_id": self.client_id,
+                        "client_secret": self.client_secret,
+                        "grant_type": "authorization_code",
+                        "code": code,
+                        "redirect_uri": self.redirect_uri,
+                    },
+                    headers={"Content-Type": "application/x-www-form-urlencoded"},
+                )
+
+                if response.status_code != 200:  # noqa: PLR2004
+                    logger.error(
+                        "Token exchange failed with provider '%s': HTTP %d",
+                        self.provider.name,
+                        response.status_code,
+                    )
+                    return None
+
+                token_data = response.json()
+                logger.debug("Token exchange successful with provider '%s'", self.provider.name)
+                return token_data.get("access_token")
+
+        except Exception as e:
+            logger.exception("Token exchange error with provider '%s': %s", self.provider.name, str(e))
+            return None
+
+    async def authenticate_with_oauth2(self, oauth_credentials: OAuth2Credentials) -> AuthenticationResponse:
+        """
+        Authenticate user with OAuth2 access token.
+
+        Args:
+            oauth_credentials: OAuth2 credentials with access token
+
+        Returns:
+            AuthenticationResponse with user and session ID
+        """
+        logger.debug("Authenticating user with OAuth2 provider '%s'", self.provider.name)
         try:
-            payload = jwt.decode(token, self.jwt_secret, algorithms=[self.jwt_algorithm])
-
-            # Reconstruct user from jwt_token payload
-            user = User(
-                user_id=payload["user_id"],
-                username=payload["username"],
-                email=payload.get("email"),
-                full_name=payload.get("full_name"),
-                roles=payload.get("roles", []),
-                metadata=payload.get("metadata", {}),
+            # Fetch user info from provider
+            async with httpx.AsyncClient() as client:
+                response = await client.get(
+                    self.provider.user_info_url,
+                    headers={"Authorization": f"{oauth_credentials.token_type} {oauth_credentials.access_token}"},
+                )
+
+                if response.status_code != 200:  # noqa: PLR2004
+                    logger.error(
+                        "Failed to fetch user info from '%s': HTTP %d",
+                        self.provider.name,
+                        response.status_code,
+                    )
+                    return AuthenticationResponse(
+                        success=False,
+                        error_message=f"Failed to fetch user info from {self.provider.name}: {response.status_code}",
+                    )
+
+                user_data = response.json()
+
+            # Create User object from provider data
+            user = self.provider.create_user_from_data(user_data)
+            logger.debug("User info fetched successfully from '%s' for user: %s", self.provider.name, user.username)
+
+            # Create session
+            now = datetime.now(timezone.utc)
+            session = Session(
+                session_id="",  # Will be generated by session store
+                user=user,
+                provider=self.provider.name,
+                oauth_token=oauth_credentials.access_token,
+                token_type=oauth_credentials.token_type,
+                created_at=now,
+                expires_at=now + timedelta(hours=self.session_expiry_hours),
+            )
+
+            session_id = await self.session_store.create_session(session)
+            logger.info(
+                "User '%s' authenticated successfully via OAuth2 provider '%s'",
+                user.username,
+                self.provider.name,
             )
 
-            return AuthenticationResponse(success=True, user=user)
+            return AuthenticationResponse(success=True, user=user, session_id=session_id)
 
-        except ExpiredSignatureError:
-            return AuthenticationResponse(success=False, error_message="Token expired")
-        except JWTError:
-            return AuthenticationResponse(success=False, error_message="Invalid jwt_token")
+        except Exception as e:
+            logger.exception("OAuth2 authentication failed with provider '%s': %s", self.provider.name, str(e))
+            return AuthenticationResponse(
+                success=False,
+                error_message=f"OAuth2 authentication failed: {str(e)}",
+            )
 
-    async def authenticate_with_oauth2(self, oauth_credentials: OAuth2Credentials) -> AuthenticationResponse:  # noqa: PLR6301
+    async def authenticate_with_credentials(self, credentials: UserCredentials) -> AuthenticationResponse:  # noqa: PLR6301
         """
-        Authenticate user with OAuth2 credentials.
+        OAuth2 backend does not support credential authentication.
 
         Args:
-            oauth_credentials: OAuth2 credentials
+            credentials: User credentials
 
         Returns:
-            AuthenticationResponse: Authentication failure as OAuth2 is not supported
+            AuthenticationResponse with error
         """
-        return AuthenticationResponse(success=False, error_message="OAuth2 not supported by ListAuthentication")
+        return AuthenticationResponse(
+            success=False,
+            error_message="Credential authentication not supported by OAuth2 backend",
+        )
 
-    async def revoke_token(self, token: str) -> bool:  # noqa: PLR6301
+    async def validate_session(self, session_id: str) -> AuthenticationResponse:
         """
-        Revoke a JWT token.
+        Validate a session.
 
         Args:
-            token: The JWT token to revoke
+            session_id: The session ID to validate
 
-        Raises:
-            NotImplementedError: This method is not implemented
+        Returns:
+            AuthenticationResponse with user if valid
         """
-        raise NotImplementedError(
-            "ListAuthenticationBackend is designed to run in development / testing scenarios. "
-            "Revoking tokens is not implemented in this backend, "
-            "if you need to revoke tokens please consider using different backend or implementing your own."
+        logger.debug(
+            "Validating OAuth2 session: %s...",
+            session_id[:SESSION_ID_LOG_LENGTH] if len(session_id) >= SESSION_ID_LOG_LENGTH else session_id,
         )
+        session = await self.session_store.get_session(session_id)
 
-    def cleanup_expired_tokens(self) -> int:
+        if not session:
+            logger.debug("OAuth2 session validation failed: session not found or expired")
+            return AuthenticationResponse(success=False, error_message="Invalid or expired session")
+
+        logger.debug("OAuth2 session validated successfully for user: %s", session.user.username)
+        return AuthenticationResponse(success=True, user=session.user)
+
+    async def revoke_session(self, session_id: str) -> bool:
         """
-        Remove expired tokens from the blacklist to prevent memory bloat.
+        Revoke a session.
+
+        Args:
+            session_id: The session ID to revoke
 
         Returns:
-            Number of tokens removed
+            True if session was revoked
         """
-        tokens_to_remove = []
+        logger.debug(
+            "Revoking OAuth2 session: %s...",
+            session_id[:SESSION_ID_LOG_LENGTH] if len(session_id) >= SESSION_ID_LOG_LENGTH else session_id,
+        )
+        success = await self.session_store.delete_session(session_id)
+        if success:
+            logger.info("OAuth2 session revoked successfully")
+        else:
+            logger.debug("OAuth2 session revocation failed: session not found")
+        return success
 
-        for token in self.revoked_tokens:
-            try:
-                # Try to decode the token - if it raises ExpiredSignatureError, it's expired
-                jwt.decode(token, self.jwt_secret, algorithms=[self.jwt_algorithm])
-            except ExpiredSignatureError:
-                # Token is expired, safe to remove from blacklist
-                tokens_to_remove.append(token)
-            except JWTError:
-                # Token is invalid, remove it too
-                tokens_to_remove.append(token)
 
-        for token in tokens_to_remove:
-            self.revoked_tokens.remove(token)
+class MultiAuthenticationBackend(AuthenticationBackend):
+    """
+    Authentication backend that supports multiple authentication methods.
 
-        return len(tokens_to_remove)
+    This backend allows combining credentials-based and OAuth2 authentication,
+    enabling users to choose their preferred login method.
+    """
+
+    def __init__(
+        self,
+        backends: list[AuthenticationBackend],
+        default_options: AuthOptions | None = None,
+    ):
+        """
+        Initialize multi-authentication backend.
+
+        Args:
+            backends: List of authentication backends to support
+            default_options: Default options for the component
+        """
+        if not backends:
+            raise ValueError("At least one authentication backend must be provided")
+
+        if default_options is None:
+            default_options = AuthOptions()
+        super().__init__(default_options)
+
+        self.backends = backends
+
+    def get_oauth2_backends(self) -> list[OAuth2AuthenticationBackend]:
+        """Get all OAuth2 backends."""
+        return [b for b in self.backends if isinstance(b, OAuth2AuthenticationBackend)]
+
+    def get_credentials_backends(self) -> list[AuthenticationBackend]:
+        """Get all credentials-based backends."""
+        return [b for b in self.backends if not isinstance(b, OAuth2AuthenticationBackend)]
+
+    async def authenticate_with_credentials(self, credentials: UserCredentials) -> AuthenticationResponse:
+        """
+        Try to authenticate with credentials using all credentials-based backends.
+
+        Args:
+            credentials: User credentials
+
+        Returns:
+            AuthenticationResponse from the first successful backend
+        """
+        logger.debug(
+            "Multi-backend credential authentication for user '%s' with %d backends",
+            credentials.username,
+            len(self.get_credentials_backends()),
+        )
+        errors = []
+
+        for backend in self.get_credentials_backends():
+            result = await backend.authenticate_with_credentials(credentials)
+            if result.success:
+                logger.debug("Credential authentication succeeded with backend: %s", type(backend).__name__)
+                return result
+            if result.error_message:
+                errors.append(result.error_message)
+
+        # All backends failed
+        error_msg = "; ".join(errors) if errors else "Authentication failed"
+        logger.warning("All credential backends failed for user '%s': %s", credentials.username, error_msg)
+        return AuthenticationResponse(success=False, error_message=error_msg)
+
+    async def authenticate_with_oauth2(self, oauth_credentials: OAuth2Credentials) -> AuthenticationResponse:
+        """
+        Try to authenticate with OAuth2 using all OAuth2 backends.
+
+        Args:
+            oauth_credentials: OAuth2 credentials
+
+        Returns:
+            AuthenticationResponse from the first successful backend
+        """
+        logger.debug("Multi-backend OAuth2 authentication with %d backends", len(self.get_oauth2_backends()))
+        errors = []
+
+        for backend in self.get_oauth2_backends():
+            result = await backend.authenticate_with_oauth2(oauth_credentials)
+            if result.success:
+                logger.debug("OAuth2 authentication succeeded with backend: %s", type(backend).__name__)
+                return result
+            if result.error_message:
+                errors.append(result.error_message)
+
+        # All backends failed
+        error_msg = "; ".join(errors) if errors else "OAuth2 authentication failed"
+        logger.warning("All OAuth2 backends failed: %s", error_msg)
+        return AuthenticationResponse(success=False, error_message=error_msg)
+
+    async def validate_session(self, session_id: str) -> AuthenticationResponse:
+        """
+        Validate a session.
+
+        Args:
+            session_id: The session ID to validate
+
+        Returns:
+            AuthenticationResponse with user if valid
+        """
+        # Try to get session from backends that have session stores
+        for backend in self.backends:
+            if hasattr(backend, "session_store") and backend.session_store:
+                result = await backend.validate_session(session_id)
+                if result.success:
+                    return result
+
+        return AuthenticationResponse(success=False, error_message="Invalid or expired session")
+
+    async def revoke_session(self, session_id: str) -> bool:
+        """
+        Revoke a session.
+
+        Args:
+            session_id: The session ID to revoke
+
+        Returns:
+            True if any backend successfully revoked the session
+        """
+        for backend in self.backends:
+            if hasattr(backend, "session_store") and backend.session_store:
+                try:
+                    success = await backend.revoke_session(session_id)
+                    if success:
+                        return True
+                except Exception:  # noqa: S112
+                    # Silently continue to next backend if one fails
+                    continue
+
+        return False