fastapi-rtk 0.2.27__py3-none-any.whl → 1.0.13__py3-none-any.whl

fastapi_rtk/routers.py

@@ -1,21 +1,29 @@
-from typing import Any, Dict, List, Optional, Tuple, Type
+import typing
 
 import jwt
+import pydantic
 from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
 from fastapi.responses import HTMLResponse, RedirectResponse
 from fastapi.security import OAuth2PasswordRequestForm
 from fastapi_users import models, schemas
 from fastapi_users.authentication import AuthenticationBackend, Authenticator, Strategy
 from fastapi_users.exceptions import UserAlreadyExists
-from fastapi_users.jwt import SecretType, decode_jwt, generate_jwt
-from fastapi_users.manager import BaseUserManager, UserManagerDependency
+from fastapi_users.jwt import SecretType, decode_jwt
+from fastapi_users.manager import UserManagerDependency
 from fastapi_users.openapi import OpenAPIResponseType
 from fastapi_users.router.common import ErrorCode, ErrorModel
-from fastapi_users.router.oauth import STATE_TOKEN_AUDIENCE
-from httpx_oauth.integrations.fastapi import OAuth2AuthorizeCallback
+from fastapi_users.router.oauth import STATE_TOKEN_AUDIENCE, generate_state_token
+from httpx_oauth.integrations.fastapi import (
+    OAuth2AuthorizeCallback as BaseOAuth2AuthorizeCallback,
+)
+from httpx_oauth.integrations.fastapi import OAuth2AuthorizeCallbackError
 from httpx_oauth.oauth2 import BaseOAuth2, OAuth2Token
 
+from .const import AuthType
+from .const import ErrorCode as _ErrorCode
 from .globals import g
+from .manager import UserManager
+from .utils import smart_run
 
 __all__ = ["get_oauth_router", "get_oauth_associate_router", "get_auth_router"]
 
@@ -32,49 +40,124 @@ HTML_SUCCESS_RESPONSE = """
 </html>
 """
 
-POPUP_DICT = {}
 
+class OAuth2AuthorizeCallback(BaseOAuth2AuthorizeCallback):
+    """
+    Modified version of the `OAuth2AuthorizeCallback` from `httpx_oauth.integrations.fastapi`.
+    - Adds `state_secret` parameter to support for decoding `state` JWT to get `redirect_url`.
+    """
+
+    state_secret: SecretType | None
+
+    def __init__(
+        self,
+        client,
+        route_name=None,
+        redirect_url=None,
+        state_secret: SecretType | None = None,
+    ):
+        self.state_secret = state_secret
+        # Use `redirect_url` from `state`
+        if self.state_secret:
+            route_name = None
+            redirect_url = ""
+        super().__init__(client, route_name, redirect_url)
+
+    async def __call__(
+        self,
+        request: Request,
+        code: typing.Optional[str] = None,
+        code_verifier: typing.Optional[str] = None,
+        state: typing.Optional[str] = None,
+        error: typing.Optional[str] = None,
+    ):
+        original_redirect_url = self.redirect_url
+        if self.state_secret:
+            # Try to decode JWT
+            try:
+                state_data = decode_jwt(
+                    state, self.state_secret, [STATE_TOKEN_AUDIENCE]
+                )
+                request.state.state_data = state_data
+                self.redirect_url = state_data.get(
+                    "redirect_url", original_redirect_url
+                )
+            except jwt.PyJWTError:
+                raise OAuth2AuthorizeCallbackError(
+                    status.HTTP_400_BAD_REQUEST,
+                    "Invalid state token",
+                )
 
-def generate_state_token(
-    data: Dict[str, str], secret: SecretType, lifetime_seconds: int = 3600
-) -> str:
-    data["aud"] = STATE_TOKEN_AUDIENCE
-    return generate_jwt(data, secret, lifetime_seconds)
+        try:
+            return await super().__call__(request, code, code_verifier, state, error)
+        finally:
+            self.redirect_url = original_redirect_url
 
 
 def get_oauth_router(
     oauth_client: BaseOAuth2,
-    backend: AuthenticationBackend,
+    backend: AuthenticationBackend[models.UP, models.ID],
     get_user_manager: UserManagerDependency[models.UP, models.ID],
     state_secret: SecretType,
-    redirect_url: Optional[str] = None,
+    redirect_url: typing.Optional[str] = None,
+    redirect_url_factory: typing.Optional[
+        typing.Callable[[Request, list[str]], str]
+    ] = None,
+    redirect_url_after_callback: typing.Optional[str] = None,
     associate_by_email: bool = False,
     is_verified_by_default: bool = False,
-    **kwargs: Dict[str, Any],
+    **kwargs: dict[str, typing.Any],
 ) -> APIRouter:
     """
-    Modified version of the original function `get_oauth_router` in `fastapi_users.router.oauth` to add:
-    - `OAUTH_REDIRECT_URI` in config is used to redirect the user to a specific URL after logging in.
-    - add `**kwargs` to accept additional arguments for the user manager's `oauth_callback` method.
+    Modified version of the `get_oauth_router` from `fastapi_users.router.oauth`.
+    - Updates the `login` to redirect to `authorization_url` directly instead of returning a JSONResponse with key `authorization_url`.
+    - Adds `redirect_url_factory` to allow dynamic generation of the redirect URL.
+    - Adds `redirect_url_after_callback` to redirect user to the application after OAuth callback.
+    - Adds `**kwargs` to provide additional arguments for the user manager's `oauth_callback` method.
 
     Generate a router with the OAuth routes.
     """
     router = APIRouter()
     callback_route_name = f"oauth:{oauth_client.name}.{backend.name}.callback"
 
-    oauth2_authorize_callback = OAuth2AuthorizeCallback(
-        oauth_client,
-        route_name=callback_route_name,
-    )
+    if redirect_url is not None:
+        oauth2_authorize_callback = OAuth2AuthorizeCallback(
+            oauth_client,
+            redirect_url=redirect_url,
+        )
+    elif redirect_url_factory is not None:
+        oauth2_authorize_callback = OAuth2AuthorizeCallback(
+            oauth_client,
+            state_secret=state_secret,
+        )
+    else:
+        oauth2_authorize_callback = OAuth2AuthorizeCallback(
+            oauth_client,
+            route_name=callback_route_name,
+        )
 
     @router.get(
         f"/login/{oauth_client.name}",
         name=f"oauth:{oauth_client.name}.{backend.name}.authorize",
     )
-    async def authorize(request: Request, scopes: List[str] = Query(None)):
-        authorize_redirect_url = str(request.url_for(callback_route_name))
+    async def authorize(request: Request, scopes: list[str] = Query(None)):
+        if redirect_url is not None:
+            authorize_redirect_url = redirect_url
+        elif redirect_url_factory is not None:
+            authorize_redirect_url = await smart_run(
+                redirect_url_factory, request, scopes
+            )
+        else:
+            authorize_redirect_url = str(request.url_for(callback_route_name))
 
-        state_data: Dict[str, str] = {}
+        state_data = dict[str, str]()
+        if redirect_url_factory:
+            state_data["redirect_url"] = authorize_redirect_url
+        if (
+            request.query_params.get("popup") == "1"
+            or request.query_params.get("popup") == "true"
+        ):
+            state_data["popup"] = "true"
         state = generate_state_token(state_data, state_secret)
         authorization_url = await oauth_client.get_authorization_url(
             authorize_redirect_url,
@@ -82,13 +165,6 @@ def get_oauth_router(
             scopes,
         )
 
-        if (
-            request.query_params.get("popup") == "1"
-            or request.query_params.get("popup") == "true"
-        ):
-            # Store the popup in POPUP_DICT
-            POPUP_DICT[state] = True
-
         return RedirectResponse(authorization_url)
 
     @router.get(
@@ -117,10 +193,10 @@ def get_oauth_router(
     )
     async def callback(
         request: Request,
-        access_token_state: Tuple[OAuth2Token, str] = Depends(
+        access_token_state: tuple[OAuth2Token, str] = Depends(
             oauth2_authorize_callback
         ),
-        user_manager: BaseUserManager[models.UP, models.ID] = Depends(get_user_manager),
+        user_manager: UserManager = Depends(get_user_manager),
         strategy: Strategy[models.UP, models.ID] = Depends(backend.get_strategy),
     ):
         token, state = access_token_state
@@ -134,10 +210,13 @@ def get_oauth_router(
                 detail=ErrorCode.OAUTH_NOT_AVAILABLE_EMAIL,
             )
 
-        try:
-            decode_jwt(state, state_secret, [STATE_TOKEN_AUDIENCE])
-        except jwt.DecodeError:
-            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
+        # Try to use existing state data from request.state
+        state_data = getattr(request.state, "state_data", None)
+        if state_data is None:
+            try:
+                state_data = decode_jwt(state, state_secret, [STATE_TOKEN_AUDIENCE])
+            except jwt.DecodeError:
+                raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
 
         try:
             user = await user_manager.oauth_callback(
@@ -169,17 +248,17 @@ def get_oauth_router(
         response = await backend.login(strategy, user)
         await user_manager.on_after_login(user, request, response)
 
-        if state in POPUP_DICT:
-            # Remove the popup from POPUP_DICT
-            POPUP_DICT.pop(state)
-
+        # Check if the request is from a popup
+        if state_data.get("popup") == "true":
             return HTMLResponse(
                 content=HTML_SUCCESS_RESPONSE,
                 headers=response.headers,
             )
 
-        if redirect_url:
-            return RedirectResponse(redirect_url, headers=response.headers)
+        if redirect_url_after_callback:
+            return RedirectResponse(
+                redirect_url_after_callback, headers=response.headers
+            )
 
         return response
 
@@ -188,17 +267,24 @@ def get_oauth_router(
 
 def get_oauth_associate_router(
     oauth_client: BaseOAuth2,
-    authenticator: Authenticator,
+    authenticator: Authenticator[models.UP, models.ID],
     get_user_manager: UserManagerDependency[models.UP, models.ID],
-    user_schema: Type[schemas.U],
+    user_schema: type[schemas.U],
     state_secret: SecretType,
-    redirect_url: Optional[str] = None,
+    redirect_url: typing.Optional[str] = None,
+    redirect_url_factory: typing.Optional[
+        typing.Callable[[Request, list[str]], str]
+    ] = None,
+    redirect_url_after_callback: typing.Optional[str] = None,
     requires_verification: bool = False,
 ) -> APIRouter:
     """
-    Generate a router with the OAuth routes to associate an authenticated user.
+    Modified version of the `get_oauth_associate_router` from `fastapi_users.router.oauth`.
+    - Updates the `login` to redirect to `authorization_url` directly instead of returning a JSONResponse with key `authorization_url`.
+    - Adds `redirect_url_factory` to allow dynamic generation of the redirect URL.
+    - Adds `redirect_url_after_callback` to redirect user to the application after OAuth callback.
 
-    Modified version of the original function in fastapi_users.router.oauth so that it redirects the user directly to the OAuth provider's login page.
+    Generate a router with the OAuth routes to associate an authenticated user.
     """
     router = APIRouter()
 
@@ -208,10 +294,21 @@ def get_oauth_associate_router(
 
     callback_route_name = f"oauth-associate:{oauth_client.name}.callback"
 
-    oauth2_authorize_callback = OAuth2AuthorizeCallback(
-        oauth_client,
-        route_name=callback_route_name,
-    )
+    if redirect_url is not None:
+        oauth2_authorize_callback = OAuth2AuthorizeCallback(
+            oauth_client,
+            redirect_url=redirect_url,
+        )
+    elif redirect_url_factory is not None:
+        oauth2_authorize_callback = OAuth2AuthorizeCallback(
+            oauth_client,
+            state_secret=state_secret,
+        )
+    else:
+        oauth2_authorize_callback = OAuth2AuthorizeCallback(
+            oauth_client,
+            route_name=callback_route_name,
+        )
 
     @router.get(
         f"/login/{oauth_client.name}",
@@ -219,12 +316,21 @@ def get_oauth_associate_router(
     )
     async def authorize(
         request: Request,
-        scopes: List[str] = Query(None),
+        scopes: list[str] = Query(None),
         user: models.UP = Depends(get_current_active_user),
     ):
-        authorize_redirect_url = str(request.url_for(callback_route_name))
+        if redirect_url is not None:
+            authorize_redirect_url = redirect_url
+        elif redirect_url_factory is not None:
+            authorize_redirect_url = await smart_run(
+                redirect_url_factory, request, scopes
+            )
+        else:
+            authorize_redirect_url = str(request.url_for(callback_route_name))
 
-        state_data: Dict[str, str] = {"sub": str(user.id)}
+        state_data = {"sub": str(user.id)}
+        if redirect_url_factory:
+            state_data["redirect_url"] = authorize_redirect_url
         state = generate_state_token(state_data, state_secret)
         authorization_url = await oauth_client.get_authorization_url(
             authorize_redirect_url,
@@ -257,10 +363,10 @@ def get_oauth_associate_router(
     async def callback(
         request: Request,
         user: models.UP = Depends(get_current_active_user),
-        access_token_state: Tuple[OAuth2Token, str] = Depends(
+        access_token_state: tuple[OAuth2Token, str] = Depends(
             oauth2_authorize_callback
         ),
-        user_manager: BaseUserManager[models.UP, models.ID] = Depends(get_user_manager),
+        user_manager: UserManager = Depends(get_user_manager),
     ):
         token, state = access_token_state
         account_id, account_email = await oauth_client.get_id_email(
@@ -273,10 +379,13 @@ def get_oauth_associate_router(
                 detail=ErrorCode.OAUTH_NOT_AVAILABLE_EMAIL,
             )
 
-        try:
-            state_data = decode_jwt(state, state_secret, [STATE_TOKEN_AUDIENCE])
-        except jwt.DecodeError:
-            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
+        # Try to use existing state data from request.state
+        state_data = getattr(request.state, "state_data", None)
+        if state_data is None:
+            try:
+                state_data = decode_jwt(state, state_secret, [STATE_TOKEN_AUDIENCE])
+            except jwt.DecodeError:
+                raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
 
         if state_data["sub"] != str(user.id):
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
@@ -292,8 +401,8 @@ def get_oauth_associate_router(
             request,
         )
 
-        if redirect_url:
-            return RedirectResponse(redirect_url)
+        if redirect_url_after_callback:
+            return RedirectResponse(redirect_url_after_callback)
 
         return schemas.model_validate(user_schema, user)
 
@@ -301,14 +410,18 @@ def get_oauth_associate_router(
 
 
 def get_auth_router(
-    backend: AuthenticationBackend,
+    backend: AuthenticationBackend[models.UP, models.ID],
     get_user_manager: UserManagerDependency[models.UP, models.ID],
-    authenticator: Authenticator,
+    authenticator: Authenticator[models.UP, models.ID],
     requires_verification: bool = False,
+    auth_type: str | None = None,
 ) -> APIRouter:
-    """Generate a router with login/logout routes for an authentication backend.
+    """
+    Modified version of the `get_auth_router` from `fastapi_users.router.auth`.
+    - Sets `g.user` to `None` after logout.
+    - Adds `auth_type` parameter to set various authentication methods like `ldap`.
 
-    Modified version of the original function in fastapi_users.router.auth so that it sets g.user to None after logout.
+    Generate a router with login/logout routes for an authentication backend.
     """
     router = APIRouter()
     get_current_user_token = authenticator.current_user_token(
@@ -344,10 +457,13 @@ def get_auth_router(
     async def login(
         request: Request,
         credentials: OAuth2PasswordRequestForm = Depends(),
-        user_manager: BaseUserManager[models.UP, models.ID] = Depends(get_user_manager),
+        user_manager: UserManager = Depends(get_user_manager),
         strategy: Strategy[models.UP, models.ID] = Depends(backend.get_strategy),
     ):
-        user = await user_manager.authenticate(credentials)
+        if auth_type == AuthType.LDAP:
+            user = await user_manager.authenticate_ldap(credentials)
+        else:
+            user = await user_manager.authenticate(credentials)
 
         if user is None or not user.is_active:
             raise HTTPException(
@@ -363,10 +479,30 @@ def get_auth_router(
         await user_manager.on_after_login(user, request, response)
         return response
 
+    class LoginBody(pydantic.BaseModel):
+        username: str
+        password: str
+
+    @router.post(
+        "/login-json",
+        name=f"auth:{backend.name}.login-json",
+        responses=login_responses,
+    )
+    async def login_with_json(
+        request: Request,
+        credentials: LoginBody,
+        user_manager: UserManager = Depends(get_user_manager),
+        strategy: Strategy[models.UP, models.ID] = Depends(backend.get_strategy),
+    ):
+        credentials = OAuth2PasswordRequestForm(
+            username=credentials.username, password=credentials.password
+        )
+        return await login(request, credentials, user_manager, strategy)
+
     logout_responses: OpenAPIResponseType = {
         **{
             status.HTTP_401_UNAUTHORIZED: {
-                "description": "Missing token or inactive user."
+                "description": _ErrorCode.GET_USER_MISSING_TOKEN_OR_INACTIVE_USER
             }
         },
         **backend.transport.get_openapi_logout_responses_success(),
@@ -377,8 +513,8 @@ def get_auth_router(
     )
     async def logout(
         request: Request,
-        user_token: Tuple[models.UP, str] = Depends(get_current_user_token),
-        user_manager: BaseUserManager[models.UP, models.ID] = Depends(get_user_manager),
+        user_token: tuple[models.UP, str] = Depends(get_current_user_token),
+        user_manager: UserManager = Depends(get_user_manager),
         strategy: Strategy[models.UP, models.ID] = Depends(backend.get_strategy),
     ):
         user, token = user_token