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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
|
from __future__ import annotations
import hashlib
import hmac
import secrets
from secrets import compare_digest
from typing import TYPE_CHECKING, Any
from litestar.datastructures import MutableScopeHeaders
from litestar.datastructures.cookie import Cookie
from litestar.enums import RequestEncodingType, ScopeType
from litestar.exceptions import PermissionDeniedException
from litestar.middleware._utils import (
build_exclude_path_pattern,
should_bypass_middleware,
)
from litestar.middleware.base import MiddlewareProtocol
from litestar.utils.scope.state import ScopeState
if TYPE_CHECKING:
from litestar.config.csrf import CSRFConfig
from litestar.connection import Request
from litestar.types import (
ASGIApp,
HTTPSendMessage,
Message,
Receive,
Scope,
Scopes,
Send,
)
__all__ = ("CSRFMiddleware",)
CSRF_SECRET_BYTES = 32
CSRF_SECRET_LENGTH = CSRF_SECRET_BYTES * 2
def generate_csrf_hash(token: str, secret: str) -> str:
"""Generate an HMAC that signs the CSRF token.
Args:
token: A hashed token.
secret: A secret value.
Returns:
A CSRF hash.
"""
return hmac.new(secret.encode(), token.encode(), hashlib.sha256).hexdigest()
def generate_csrf_token(secret: str) -> str:
"""Generate a CSRF token that includes a randomly generated string signed by an HMAC.
Args:
secret: A secret string.
Returns:
A unique CSRF token.
"""
token = secrets.token_hex(CSRF_SECRET_BYTES)
token_hash = generate_csrf_hash(token=token, secret=secret)
return token + token_hash
class CSRFMiddleware(MiddlewareProtocol):
"""CSRF Middleware class.
This Middleware protects against attacks by setting a CSRF cookie with a token and verifying it in request headers.
"""
scopes: Scopes = {ScopeType.HTTP}
def __init__(self, app: ASGIApp, config: CSRFConfig) -> None:
"""Initialize ``CSRFMiddleware``.
Args:
app: The ``next`` ASGI app to call.
config: The CSRFConfig instance.
"""
self.app = app
self.config = config
self.exclude = build_exclude_path_pattern(exclude=config.exclude)
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
"""ASGI callable.
Args:
scope: The ASGI connection scope.
receive: The ASGI receive function.
send: The ASGI send function.
Returns:
None
"""
if scope["type"] != ScopeType.HTTP:
await self.app(scope, receive, send)
return
request: Request[Any, Any, Any] = scope["app"].request_class(scope=scope, receive=receive)
content_type, _ = request.content_type
csrf_cookie = request.cookies.get(self.config.cookie_name)
existing_csrf_token = request.headers.get(self.config.header_name)
if not existing_csrf_token and content_type in {
RequestEncodingType.URL_ENCODED,
RequestEncodingType.MULTI_PART,
}:
form = await request.form()
existing_csrf_token = form.get("_csrf_token", None)
connection_state = ScopeState.from_scope(scope)
if request.method in self.config.safe_methods or should_bypass_middleware(
scope=scope,
scopes=self.scopes,
exclude_opt_key=self.config.exclude_from_csrf_key,
exclude_path_pattern=self.exclude,
):
token = connection_state.csrf_token = csrf_cookie or generate_csrf_token(secret=self.config.secret)
await self.app(scope, receive, self.create_send_wrapper(send=send, csrf_cookie=csrf_cookie, token=token))
elif (
existing_csrf_token is not None
and csrf_cookie is not None
and self._csrf_tokens_match(existing_csrf_token, csrf_cookie)
):
connection_state.csrf_token = existing_csrf_token
await self.app(scope, receive, send)
else:
raise PermissionDeniedException("CSRF token verification failed")
def create_send_wrapper(self, send: Send, token: str, csrf_cookie: str | None) -> Send:
"""Wrap ``send`` to handle CSRF validation.
Args:
token: The CSRF token.
send: The ASGI send function.
csrf_cookie: CSRF cookie.
Returns:
An ASGI send function.
"""
async def send_wrapper(message: Message) -> None:
"""Send function that wraps the original send to inject a cookie.
Args:
message: An ASGI ``Message``
Returns:
None
"""
if csrf_cookie is None and message["type"] == "http.response.start":
message.setdefault("headers", [])
self._set_cookie_if_needed(message=message, token=token)
await send(message)
return send_wrapper
def _set_cookie_if_needed(self, message: HTTPSendMessage, token: str) -> None:
headers = MutableScopeHeaders.from_message(message)
cookie = Cookie(
key=self.config.cookie_name,
value=token,
path=self.config.cookie_path,
secure=self.config.cookie_secure,
httponly=self.config.cookie_httponly,
samesite=self.config.cookie_samesite,
domain=self.config.cookie_domain,
)
headers.add("set-cookie", cookie.to_header(header=""))
def _decode_csrf_token(self, token: str) -> str | None:
"""Decode a CSRF token and validate its HMAC."""
if len(token) < CSRF_SECRET_LENGTH + 1:
return None
token_secret = token[:CSRF_SECRET_LENGTH]
existing_hash = token[CSRF_SECRET_LENGTH:]
expected_hash = generate_csrf_hash(token=token_secret, secret=self.config.secret)
return token_secret if compare_digest(existing_hash, expected_hash) else None
def _csrf_tokens_match(self, request_csrf_token: str, cookie_csrf_token: str) -> bool:
"""Take the CSRF tokens from the request and the cookie and verify both are valid and identical."""
decoded_request_token = self._decode_csrf_token(request_csrf_token)
decoded_cookie_token = self._decode_csrf_token(cookie_csrf_token)
if decoded_request_token is None or decoded_cookie_token is None:
return False
return compare_digest(decoded_request_token, decoded_cookie_token)
|
