|
22 | 22 | ) |
23 | 23 | from utils.core.auth import refresh_token_is_persistent, set_auth_cookies |
24 | 24 | from utils.core.rate_limit import get_trusted_proxy_hosts |
| 25 | +from utils.core.csrf import ( |
| 26 | + CSRF_COOKIE_NAME, |
| 27 | + UNSAFE_HTTP_METHODS, |
| 28 | + csrf_enabled, |
| 29 | + extract_submitted_csrf_token, |
| 30 | + generate_csrf_token, |
| 31 | + set_csrf_cookie, |
| 32 | + validate_csrf_token, |
| 33 | +) |
25 | 34 | from utils.core.htmx import ( |
26 | 35 | is_htmx_request, |
27 | 36 | toast_response, |
|
31 | 40 | from exceptions.http_exceptions import ( |
32 | 41 | AlreadyAuthenticatedError, |
33 | 42 | AuthenticationError, |
| 43 | + CsrfError, |
34 | 44 | PasswordValidationError, |
35 | 45 | CredentialsError, |
36 | 46 | RateLimitError, |
@@ -80,6 +90,22 @@ async def flash_cookie_middleware(request: Request, call_next): |
80 | 90 | return response |
81 | 91 |
|
82 | 92 |
|
| 93 | +@app.middleware("http") |
| 94 | +async def csrf_middleware(request: Request, call_next): |
| 95 | + token = request.cookies.get(CSRF_COOKIE_NAME) or generate_csrf_token() |
| 96 | + request.state.csrf_token = token |
| 97 | + |
| 98 | + if csrf_enabled() and request.method in UNSAFE_HTTP_METHODS: |
| 99 | + submitted = await extract_submitted_csrf_token(request) |
| 100 | + if not validate_csrf_token(request, submitted): |
| 101 | + return await csrf_error_handler(request, CsrfError()) |
| 102 | + |
| 103 | + response = await call_next(request) |
| 104 | + if request.cookies.get(CSRF_COOKIE_NAME) != token: |
| 105 | + set_csrf_cookie(response, token) |
| 106 | + return response |
| 107 | + |
| 108 | + |
83 | 109 | # --- Include Routers --- |
84 | 110 |
|
85 | 111 |
|
@@ -144,6 +170,25 @@ async def rate_limit_error_handler(request: Request, exc: RateLimitError): |
144 | 170 | return response |
145 | 171 |
|
146 | 172 |
|
| 173 | +@app.exception_handler(CsrfError) |
| 174 | +async def csrf_error_handler(request: Request, exc: CsrfError): |
| 175 | + if is_htmx_request(request): |
| 176 | + return toast_response( |
| 177 | + request, |
| 178 | + templates, |
| 179 | + exc.detail, |
| 180 | + level="danger", |
| 181 | + status_code=403, |
| 182 | + ) |
| 183 | + user = await get_user_from_request(request) |
| 184 | + return templates.TemplateResponse( |
| 185 | + request, |
| 186 | + "errors/error.html", |
| 187 | + {"status_code": 403, "detail": exc.detail, "errors": None, "user": user}, |
| 188 | + status_code=403, |
| 189 | + ) |
| 190 | + |
| 191 | + |
147 | 192 | # Handle CredentialsError (invalid email/password) with toast for HTMX |
148 | 193 | @app.exception_handler(CredentialsError) |
149 | 194 | async def credentials_exception_handler(request: Request, exc: CredentialsError): |
|
0 commit comments