Close #3: forgot/reset password routes

fastapi_users/authentication/jwt.py

@@ -1,5 +1,3 @@
-from datetime import datetime, timedelta
-
 import jwt
 from fastapi import Depends
 from fastapi.security import OAuth2PasswordBearer
@@ -8,17 +6,11 @@ from starlette.responses import Response
 from fastapi_users.authentication.base import BaseAuthentication
 from fastapi_users.db.base import BaseUserDatabase
 from fastapi_users.models import BaseUserDB
+from fastapi_users.utils import JWT_ALGORITHM, generate_jwt
 
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/login")
 
 
-def generate_jwt(data: dict, lifetime_seconds: int, secret: str, algorithm: str) -> str:
-    payload = data.copy()
-    expire = datetime.utcnow() + timedelta(seconds=lifetime_seconds)
-    payload["exp"] = expire
-    return jwt.encode(payload, secret, algorithm=algorithm).decode("utf-8")
-
-
 class JWTAuthentication(BaseAuthentication):
     """
     Authentication using a JWT.
@@ -27,7 +19,7 @@ class JWTAuthentication(BaseAuthentication):
     :param lifetime_seconds: Lifetime duration of the JWT in seconds.
     """
 
-    algorithm: str = "HS256"
+    token_audience: str = "fastapi-users:auth"
     secret: str
     lifetime_seconds: int
 
@@ -36,8 +28,8 @@ class JWTAuthentication(BaseAuthentication):
         self.lifetime_seconds = lifetime_seconds
 
     async def get_login_response(self, user: BaseUserDB, response: Response):
-        data = {"user_id": user.id}
-        token = generate_jwt(data, self.lifetime_seconds, self.secret, self.algorithm)
+        data = {"user_id": user.id, "aud": self.token_audience}
+        token = generate_jwt(data, self.lifetime_seconds, self.secret, JWT_ALGORITHM)
 
         return {"token": token}
 
@@ -65,7 +57,12 @@ class JWTAuthentication(BaseAuthentication):
     def _get_authentication_method(self, user_db: BaseUserDatabase):
         async def authentication_method(token: str = Depends(oauth2_scheme)):
             try:
-                data = jwt.decode(token, self.secret, algorithms=[self.algorithm])
+                data = jwt.decode(
+                    token,
+                    self.secret,
+                    audience=self.token_audience,
+                    algorithms=[JWT_ALGORITHM],
+                )
                 user_id = data.get("user_id")
                 if user_id is None:
                     return None

fastapi_users/fastapi_users.py

@@ -1,6 +1,6 @@
 """Ready-to-use and customizable users management for FastAPI."""
 
-from typing import Callable, Type
+from typing import Any, Callable, Type
 
 from fastapi import APIRouter
 
@@ -17,6 +17,9 @@ class FastAPIUsers:
     :param db: Database adapter instance.
     :param auth: Authentication logic instance.
     :param user_model: Pydantic model of a user.
+    :param on_after_forgot_password: Hook called after a forgot password request.
+    :param reset_password_token_secret: Secret to encode reset password token.
+    :param reset_password_token_lifetime_seconds: Lifetime of reset password token.
 
     :attribute router: FastAPI router exposing authentication routes.
     :attribute get_current_user: Dependency callable to inject authenticated user.
@@ -28,11 +31,24 @@ class FastAPIUsers:
     get_current_user: Callable[..., BaseUserDB]
 
     def __init__(
-        self, db: BaseUserDatabase, auth: BaseAuthentication, user_model: Type[BaseUser]
+        self,
+        db: BaseUserDatabase,
+        auth: BaseAuthentication,
+        user_model: Type[BaseUser],
+        on_after_forgot_password: Callable[[BaseUserDB, str], Any],
+        reset_password_token_secret: str,
+        reset_password_token_lifetime_seconds: int = 3600,
     ):
         self.db = db
         self.auth = auth
-        self.router = get_user_router(self.db, user_model, self.auth)
+        self.router = get_user_router(
+            self.db,
+            user_model,
+            self.auth,
+            on_after_forgot_password,
+            reset_password_token_secret,
+            reset_password_token_lifetime_seconds,
+        )
 
         get_current_user = self.auth.get_current_user(self.db)
         self.get_current_user = get_current_user  # type: ignore

fastapi_users/router.py

@@ -1,23 +1,37 @@
-from typing import Type
+import inspect
+from typing import Any, Callable, Type
 
-from fastapi import APIRouter, Depends, HTTPException
+import jwt
+from fastapi import APIRouter, Body, Depends, HTTPException
 from fastapi.security import OAuth2PasswordRequestForm
+from pydantic.types import EmailStr
 from starlette import status
 from starlette.responses import Response
 
 from fastapi_users.authentication import BaseAuthentication
 from fastapi_users.db import BaseUserDatabase
-from fastapi_users.models import BaseUser, Models
+from fastapi_users.models import BaseUser, BaseUserDB, Models
 from fastapi_users.password import get_password_hash
+from fastapi_users.utils import JWT_ALGORITHM, generate_jwt
 
 
 def get_user_router(
-    user_db: BaseUserDatabase, user_model: Type[BaseUser], auth: BaseAuthentication
+    user_db: BaseUserDatabase,
+    user_model: Type[BaseUser],
+    auth: BaseAuthentication,
+    on_after_forgot_password: Callable[[BaseUserDB, str], Any],
+    reset_password_token_secret: str,
+    reset_password_token_lifetime_seconds: int = 3600,
 ) -> APIRouter:
     """Generate a router with the authentication routes."""
     router = APIRouter()
     models = Models(user_model)
 
+    reset_password_token_audience = "fastapi-users:reset"
+    is_on_after_forgot_password_async = inspect.iscoroutinefunction(
+        on_after_forgot_password
+    )
+
     @router.post("/register", response_model=models.User)
     async def register(user: models.UserCreate):  # type: ignore
         hashed_password = get_password_hash(user.password)
@@ -38,4 +52,45 @@ def get_user_router(
 
         return await auth.get_login_response(user, response)
 
+    @router.post("/forgot-password", status_code=status.HTTP_202_ACCEPTED)
+    async def forgot_password(email: EmailStr = Body(..., embed=True)):
+        user = await user_db.get_by_email(email)
+
+        if user is not None and user.is_active:
+            token_data = {"user_id": user.id, "aud": reset_password_token_audience}
+            token = generate_jwt(
+                token_data,
+                reset_password_token_lifetime_seconds,
+                reset_password_token_secret,
+            )
+            if is_on_after_forgot_password_async:
+                await on_after_forgot_password(user, token)
+            else:
+                on_after_forgot_password(user, token)
+
+        return None
+
+    @router.post("/reset-password")
+    async def reset_password(token: str = Body(...), password: str = Body(...)):
+        try:
+            data = jwt.decode(
+                token,
+                reset_password_token_secret,
+                audience=reset_password_token_audience,
+                algorithms=[JWT_ALGORITHM],
+            )
+            user_id = data.get("user_id")
+            if user_id is None:
+                raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
+
+            user = await user_db.get(user_id)
+            if user is None or not user.is_active:
+                raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
+
+            updated_user = BaseUserDB(**user.dict())
+            updated_user.hashed_password = get_password_hash(password)
+            await user_db.update(updated_user)
+        except jwt.PyJWTError:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
+
     return router

fastapi_users/utils.py

@@ -0,0 +1,14 @@
+from datetime import datetime, timedelta
+
+import jwt
+
+JWT_ALGORITHM = "HS256"
+
+
+def generate_jwt(
+    data: dict, lifetime_seconds: int, secret: str, algorithm: str = JWT_ALGORITHM
+) -> str:
+    payload = data.copy()
+    expire = datetime.utcnow() + timedelta(seconds=lifetime_seconds)
+    payload["exp"] = expire
+    return jwt.encode(payload, secret, algorithm=algorithm).decode("utf-8")