sessions.py 3.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283
  1. import json
  2. import typing
  3. from base64 import b64decode, b64encode
  4. import itsdangerous
  5. from itsdangerous.exc import BadSignature
  6. from starlette.datastructures import MutableHeaders, Secret
  7. from starlette.requests import HTTPConnection
  8. from starlette.types import ASGIApp, Message, Receive, Scope, Send
  9. class SessionMiddleware:
  10. def __init__(
  11. self,
  12. app: ASGIApp,
  13. secret_key: typing.Union[str, Secret],
  14. session_cookie: str = "session",
  15. max_age: typing.Optional[int] = 14 * 24 * 60 * 60, # 14 days, in seconds
  16. path: str = "/",
  17. same_site: typing.Literal["lax", "strict", "none"] = "lax",
  18. https_only: bool = False,
  19. domain: typing.Optional[str] = None,
  20. ) -> None:
  21. self.app = app
  22. self.signer = itsdangerous.TimestampSigner(str(secret_key))
  23. self.session_cookie = session_cookie
  24. self.max_age = max_age
  25. self.path = path
  26. self.security_flags = "httponly; samesite=" + same_site
  27. if https_only: # Secure flag can be used with HTTPS only
  28. self.security_flags += "; secure"
  29. if domain is not None:
  30. self.security_flags += f"; domain={domain}"
  31. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  32. if scope["type"] not in ("http", "websocket"): # pragma: no cover
  33. await self.app(scope, receive, send)
  34. return
  35. connection = HTTPConnection(scope)
  36. initial_session_was_empty = True
  37. if self.session_cookie in connection.cookies:
  38. data = connection.cookies[self.session_cookie].encode("utf-8")
  39. try:
  40. data = self.signer.unsign(data, max_age=self.max_age)
  41. scope["session"] = json.loads(b64decode(data))
  42. initial_session_was_empty = False
  43. except BadSignature:
  44. scope["session"] = {}
  45. else:
  46. scope["session"] = {}
  47. async def send_wrapper(message: Message) -> None:
  48. if message["type"] == "http.response.start":
  49. if scope["session"]:
  50. # We have session data to persist.
  51. data = b64encode(json.dumps(scope["session"]).encode("utf-8"))
  52. data = self.signer.sign(data)
  53. headers = MutableHeaders(scope=message)
  54. header_value = "{session_cookie}={data}; path={path}; {max_age}{security_flags}".format( # noqa E501
  55. session_cookie=self.session_cookie,
  56. data=data.decode("utf-8"),
  57. path=self.path,
  58. max_age=f"Max-Age={self.max_age}; " if self.max_age else "",
  59. security_flags=self.security_flags,
  60. )
  61. headers.append("Set-Cookie", header_value)
  62. elif not initial_session_was_empty:
  63. # The session has been cleared.
  64. headers = MutableHeaders(scope=message)
  65. header_value = "{session_cookie}={data}; path={path}; {expires}{security_flags}".format( # noqa E501
  66. session_cookie=self.session_cookie,
  67. data="null",
  68. path=self.path,
  69. expires="expires=Thu, 01 Jan 1970 00:00:00 GMT; ",
  70. security_flags=self.security_flags,
  71. )
  72. headers.append("Set-Cookie", header_value)
  73. await send(message)
  74. await self.app(scope, receive, send_wrapper)