1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
|
import secrets
from os import environ
from typing import Any, Dict, Optional
from uuid import UUID
from pydantic import BaseModel, EmailStr
from litestar import Litestar, Request, Response, get, post
from litestar.connection import ASGIConnection
from litestar.openapi.config import OpenAPIConfig
from litestar.security.jwt import JWTCookieAuth, Token
# Let's assume we have a User model that is a pydantic model.
# This though is not required - we need some sort of user class -
# but it can be any arbitrary value, e.g. an SQLAlchemy model, a representation of a MongoDB etc.
class User(BaseModel):
id: UUID
name: str
email: EmailStr
MOCK_DB: Dict[str, User] = {}
# JWTCookieAuth requires a retrieve handler callable that receives the JWT token model and the ASGI connection
# and returns the 'User' instance correlating to it.
#
# Notes:
# - 'User' can be any arbitrary value you decide upon.
# - The callable can be either sync or async - both will work.
async def retrieve_user_handler(token: "Token", connection: "ASGIConnection[Any, Any, Any, Any]") -> Optional[User]:
# logic here to retrieve the user instance
return MOCK_DB.get(token.sub)
jwt_cookie_auth = JWTCookieAuth[User](
retrieve_user_handler=retrieve_user_handler,
token_secret=environ.get("JWT_SECRET", secrets.token_hex()),
# we are specifying which endpoints should be excluded from authentication. In this case the login endpoint
# and our openAPI docs.
exclude=["/login", "/schema"],
# Tip: We can optionally supply cookie options to the configuration. Here is an example of enabling the secure cookie option
# secure=True,
)
# Given an instance of 'JWTCookieAuth' we can create a login handler function:
@post("/login")
async def login_handler(data: "User") -> "Response[User]":
MOCK_DB[str(data.id)] = data
return jwt_cookie_auth.login(identifier=str(data.id), response_body=data)
# We also have some other routes, for example:
@get("/some-path", sync_to_thread=False)
def some_route_handler(request: "Request[User, Token, Any]") -> Any:
# request.user is set to the instance of user returned by the middleware
assert isinstance(request.user, User)
# request.auth is the instance of 'litestar_jwt.Token' created from the data encoded in the auth header
assert isinstance(request.auth, Token)
# do stuff ...
# We create our OpenAPIConfig as usual - the JWT security scheme will be injected into it.
openapi_config = OpenAPIConfig(
title="My API",
version="1.0.0",
)
# We initialize the app instance and pass the jwt_cookie_auth 'on_app_init' handler to the constructor.
# The hook handler will inject the JWT middleware and openapi configuration into the app.
app = Litestar(
route_handlers=[login_handler, some_route_handler],
on_app_init=[jwt_cookie_auth.on_app_init],
openapi_config=openapi_config,
)
|
