Fixed middleware error handling and handled form and password validation via Pydantic with 422 status code

Author: chriscarrollsmith

main.py

@@ -1,18 +1,21 @@
+# ToDo: Add CSRF protection to all POST, download, and sensitive data routes
+
 import logging
 from typing import Optional
 from contextlib import asynccontextmanager
 from fastapi import FastAPI, Request, Depends, status
 from fastapi.responses import RedirectResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
-from fastapi.exceptions import RequestValidationError, StarletteHTTPException, HTTPException
+from fastapi.exceptions import RequestValidationError, HTTPException, StarletteHTTPException
 from sqlmodel import Session
 from routers import authentication, organization, role, user
-from utils.auth import get_authenticated_user, get_optional_user, NeedsNewTokens, get_user_from_reset_token
+from utils.auth import get_authenticated_user, get_optional_user, NeedsNewTokens, get_user_from_reset_token, PasswordValidationError
 from utils.db import User, get_session
 
 
 logger = logging.getLogger("uvicorn.error")
+logger.setLevel(logging.DEBUG)
 
 
 @asynccontextmanager
@@ -31,7 +34,10 @@ async def lifespan(app: FastAPI):
 templates = Jinja2Templates(directory="templates")
 
 
-# Middleware to handle the NeedsNewTokens exception
+# -- Exception Handling Middlewares --
+
+
+# Handle NeedsNewTokens by setting new tokens and redirecting to same page
 @app.exception_handler(NeedsNewTokens)
 async def needs_new_tokens_handler(request: Request, exc: NeedsNewTokens):
     response = RedirectResponse(
@@ -52,28 +58,83 @@ async def needs_new_tokens_handler(request: Request, exc: NeedsNewTokens):
     )
     return response
 
-# TODO: Make sure this only catches server errors and not 307 redirects
-# Create a custom server error class that inherits from StarletteHTTPException?
-# @app.exception_handler(StarletteHTTPException)
-# async def http_exception_handler(request: Request, exc: StarletteHTTPException):
-#     return templates.TemplateResponse(
-#         "errors/error.html",
-#         {"request": request, "status_code": exc.status_code, "detail": exc.detail},
-#         status_code=exc.status_code,
-#     )
+
+# Handle PasswordValidationError by rendering the error page
+@app.exception_handler(PasswordValidationError)
+async def password_validation_exception_handler(request: Request, exc: PasswordValidationError):
+    return templates.TemplateResponse(
+        "errors/validation_error.html",
+        {
+            "request": request,
+            "status_code": 422,
+            "errors": {exc.detail["field"]: exc.detail["message"]}
+        },
+        status_code=422,
+    )
 
 
+# Handle RequestValidationError by rendering the error page (TODO: use toast instead?)
 @app.exception_handler(RequestValidationError)
 async def validation_exception_handler(request: Request, exc: RequestValidationError):
+    errors = {}
+    for error in exc.errors():
+        # Handle different error locations more carefully
+        location = error["loc"]
+
+        # Skip type errors for the whole body
+        if len(location) == 1 and location[0] == "body":
+            continue
+
+        # For form fields, the location might be just (field_name,)
+        # For JSON body, it might be (body, field_name)
+        field_name = location[-1]  # Take the last item in the location tuple
+        errors[field_name] = error["msg"]
+
     return templates.TemplateResponse(
-        "errors/error.html",
-        {"request": request, "status_code": 422, "detail": str(exc)},
+        "errors/validation_error.html",
+        {
+            "request": request,
+            "status_code": 422,
+            "errors": errors
+        },
         status_code=422,
     )
 
 
+# Handle StarletteHTTPException (including 404, 405, etc.) by rendering the error page
+@app.exception_handler(StarletteHTTPException)
+async def http_exception_handler(request: Request, exc: StarletteHTTPException):
+    # Don't handle redirects
+    if exc.status_code in [301, 302, 303, 307, 308]:
+        raise exc
+
+    return templates.TemplateResponse(
+        "errors/error.html",
+        {"request": request, "status_code": exc.status_code, "detail": exc.detail},
+        status_code=exc.status_code,
+    )
+
+
+# Add handler for uncaught exceptions (500 Internal Server Error)
+@app.exception_handler(Exception)
+async def general_exception_handler(request: Request, exc: Exception):
+    # Log the error for debugging
+    logger.error(f"Unhandled exception: {exc}", exc_info=True)
+
+    return templates.TemplateResponse(
+        "errors/error.html",
+        {
+            "request": request,
+            "status_code": 500,
+            "detail": "Internal Server Error"
+        },
+        status_code=500,
+    )
+
+
 # -- Unauthenticated Routes --
 
+
 # Define a dependency for common parameters
 async def common_unauthenticated_parameters(
     request: Request,
@@ -113,7 +174,7 @@ async def read_register(
 @app.get("/forgot_password")
 async def read_forgot_password(
     params: dict = Depends(common_unauthenticated_parameters),
-    show_form: Optional[bool] = True,
+    show_form: Optional[str] = "true",
 ):
     if params["user"]:
         return RedirectResponse(url="/dashboard", status_code=302)
@@ -196,7 +257,6 @@ async def read_profile(
 app.include_router(role.router)
 app.include_router(user.router)
 
-
 if __name__ == "__main__":
     import uvicorn
 

routers/authentication.py

@@ -1,15 +1,17 @@
 # auth.py
+from logging import getLogger
 from typing import Optional
 from datetime import datetime
-from fastapi import APIRouter, Depends, HTTPException, Form, BackgroundTasks
+from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks, Form
 from fastapi.responses import RedirectResponse
-from pydantic import BaseModel, EmailStr, ConfigDict
+from pydantic import BaseModel, EmailStr, ConfigDict, field_validator
 from sqlmodel import Session, select
 from utils.db import User
 from utils.auth import (
     get_session,
-    validate_password_strength,
     get_user_from_reset_token,
+    create_password_validator,
+    create_passwords_match_validator,
     oauth2_scheme_cookie,
     get_password_hash,
     verify_password,
@@ -19,14 +21,88 @@
     send_reset_email
 )
 
+logger = getLogger("uvicorn.error")
+
 router = APIRouter(prefix="/auth", tags=["auth"])
 
 
-class UserCreate(BaseModel):
+# -- Server Request and Response Models --
+
+
+class UserRegister(BaseModel):
     name: str
     email: EmailStr
     password: str
-    organization_id: Optional[int] = None
+    confirm_password: str
+
+    validate_password_strength = create_password_validator("password")
+    validate_passwords_match = create_passwords_match_validator(
+        "password", "confirm_password")
+
+    @classmethod
+    async def as_form(
+        cls,
+        name: str = Form(...),
+        email: EmailStr = Form(...),
+        password: str = Form(...),
+        confirm_password: str = Form(...)
+    ):
+        return cls(
+            name=name,
+            email=email,
+            password=password,
+            confirm_password=confirm_password
+        )
+
+
+class UserLogin(BaseModel):
+    email: EmailStr
+    password: str
+
+    @classmethod
+    async def as_form(
+        cls,
+        email: EmailStr = Form(...),
+        password: str = Form(...)
+    ):
+        return cls(email=email, password=password)
+
+
+class UserForgotPassword(BaseModel):
+    email: EmailStr
+
+    @classmethod
+    async def as_form(
+        cls,
+        email: EmailStr = Form(...)
+    ):
+        return cls(email=email)
+
+
+class UserResetPassword(BaseModel):
+    email: EmailStr
+    token: str
+    new_password: str
+    confirm_new_password: str
+
+    # Use the factory with a different field name
+    validate_password_strength = create_password_validator("new_password")
+    validate_passwords_match = create_passwords_match_validator(
+        "new_password", "confirm_new_password")
+
+    @classmethod
+    async def as_form(
+        cls,
+        email: EmailStr = Form(...),
+        token: str = Form(...),
+        new_password: str = Form(...),
+        confirm_new_password: str = Form(...)
+    ):
+        return cls(email=email, token=token,
+                   new_password=new_password, confirm_new_password=confirm_new_password)
+
+
+# -- DB Request and Response Models --
 
 
 class UserRead(BaseModel):
@@ -41,31 +117,14 @@ class UserRead(BaseModel):
     deleted: bool
 
 
-class UserForgotPassword(BaseModel):
-    email: EmailStr
-
-
-class UserResetPassword(BaseModel):
-    token: str
-    new_password: str
+# -- Routes --
 
 
 @router.post("/register", response_class=RedirectResponse)
 async def register(
-    name: str = Form(...),
-    email: EmailStr = Form(...),
-    password: str = Form(...),
-    confirm_password: str = Form(...),
+    user: UserRegister = Depends(UserRegister.as_form),
     session: Session = Depends(get_session),
 ) -> RedirectResponse:
-    if password != confirm_password:
-        raise HTTPException(status_code=400, detail="Passwords do not match")
-
-    if not validate_password_strength(password):
-        raise HTTPException(
-            status_code=400, detail="Password does not satisfy the security policy")
-
-    user = UserCreate(name=name, email=email, password=password)
     db_user = session.exec(select(User).where(
         User.email == user.email)).first()
 
@@ -92,13 +151,13 @@ async def register(
 
 
 @router.post("/login", response_class=RedirectResponse)
-def login(
-    email: str = Form(...),
-    password: str = Form(...),
+async def login(
+    user: UserLogin = Depends(UserLogin.as_form),
     session: Session = Depends(get_session),
 ) -> RedirectResponse:
-    db_user = session.exec(select(User).where(User.email == email)).first()
-    if not db_user or not verify_password(password, db_user.hashed_password):
+    db_user = session.exec(select(User).where(
+        User.email == user.email)).first()
+    if not db_user or not verify_password(user.password, db_user.hashed_password):
         raise HTTPException(status_code=400, detail="Invalid credentials")
 
     # Create access token
@@ -128,7 +187,7 @@ def login(
 
 # Updated refresh_token endpoint
 @router.post("/refresh", response_class=RedirectResponse)
-def refresh_token(
+async def refresh_token(
     tokens: tuple[Optional[str], Optional[str]
                   ] = Depends(oauth2_scheme_cookie),
     session: Session = Depends(get_session),
@@ -173,46 +232,35 @@ def refresh_token(
     return response
 
 
-class EmailSchema(BaseModel):
-    email: EmailStr
-
-
-class ResetSchema(BaseModel):
-    token: str
-    new_password: str
-
-
 @router.post("/forgot_password")
-def forgot_password(user: UserForgotPassword, background_tasks: BackgroundTasks, session: Session = Depends(get_session)):
+async def forgot_password(
+    background_tasks: BackgroundTasks,
+    user: UserForgotPassword = Depends(UserForgotPassword.as_form),
+    session: Session = Depends(get_session)
+):
     db_user = session.exec(select(User).where(
         User.email == user.email)).first()
 
     # TODO: Handle this in background task so we don't leak information via timing attacks
     if db_user:
         background_tasks.add_task(send_reset_email, user.email, session)
 
-    return RedirectResponse(url="/forgot_password", status_code=303, show_form=False)
+    return RedirectResponse(url="/forgot_password?show_form=false", status_code=303)
 
 
 @router.post("/reset_password")
-def reset_password(
-    email: str, token: str, new_password: str, confirm_new_password: str, session: Session = Depends(get_session)
+async def reset_password(
+    user: UserResetPassword = Depends(UserResetPassword.as_form),
+    session: Session = Depends(get_session)
 ):
-    if new_password != confirm_new_password:
-        raise HTTPException(status_code=400, detail="Passwords do not match")
-
-    if not validate_password_strength(new_password):
-        raise HTTPException(
-            status_code=400, detail="Password does not satisfy the security policy")
-
     authorized_user, reset_token = get_user_from_reset_token(
-        email, token, session)
+        user.email, user.token, session)
 
     if not authorized_user:
         raise HTTPException(status_code=400, detail="Invalid or expired token")
 
     # Update password and mark token as used
-    authorized_user.hashed_password = get_password_hash(new_password)
+    authorized_user.hashed_password = get_password_hash(user.new_password)
     reset_token.used = True
     session.commit()
     session.refresh(authorized_user)