sessions.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125
  1. from __future__ import annotations
  2. import json
  3. import typing
  4. from base64 import b64decode, b64encode
  5. from typing import Literal
  6. import itsdangerous
  7. from itsdangerous.exc import BadSignature
  8. from starlette.datastructures import MutableHeaders, Secret
  9. from starlette.requests import HTTPConnection
  10. from starlette.types import ASGIApp, Message, Receive, Scope, Send
  11. class SessionMiddleware:
  12. def __init__(
  13. self,
  14. app: ASGIApp,
  15. secret_key: str | Secret,
  16. session_cookie: str = "session",
  17. max_age: int | None = 14 * 24 * 60 * 60, # 14 days, in seconds
  18. path: str = "/",
  19. same_site: Literal["lax", "strict", "none"] = "lax",
  20. https_only: bool = False,
  21. domain: str | None = None,
  22. ) -> None:
  23. self.app = app
  24. self.signer = itsdangerous.TimestampSigner(str(secret_key))
  25. self.session_cookie = session_cookie
  26. self.max_age = max_age
  27. self.path = path
  28. self.security_flags = "httponly; samesite=" + same_site
  29. if https_only: # Secure flag can be used with HTTPS only
  30. self.security_flags += "; secure"
  31. if domain is not None:
  32. self.security_flags += f"; domain={domain}"
  33. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  34. if scope["type"] not in ("http", "websocket"): # pragma: no cover
  35. await self.app(scope, receive, send)
  36. return
  37. connection = HTTPConnection(scope)
  38. initial_session_was_empty = True
  39. if self.session_cookie in connection.cookies:
  40. data = connection.cookies[self.session_cookie].encode("utf-8")
  41. try:
  42. data = self.signer.unsign(data, max_age=self.max_age)
  43. scope["session"] = Session(json.loads(b64decode(data)))
  44. initial_session_was_empty = False
  45. except BadSignature:
  46. scope["session"] = Session()
  47. else:
  48. scope["session"] = Session()
  49. async def send_wrapper(message: Message) -> None:
  50. if message["type"] == "http.response.start":
  51. session: Session = scope["session"]
  52. headers = MutableHeaders(scope=message)
  53. if session.accessed:
  54. headers.add_vary_header("Cookie")
  55. if session.modified and session:
  56. # We have session data to persist.
  57. data = b64encode(json.dumps(session).encode("utf-8"))
  58. data = self.signer.sign(data)
  59. header_value = "{session_cookie}={data}; path={path}; {max_age}{security_flags}".format(
  60. session_cookie=self.session_cookie,
  61. data=data.decode("utf-8"),
  62. path=self.path,
  63. max_age=f"Max-Age={self.max_age}; " if self.max_age else "",
  64. security_flags=self.security_flags,
  65. )
  66. headers.append("Set-Cookie", header_value)
  67. elif session.modified and not initial_session_was_empty:
  68. # The session has been cleared.
  69. header_value = "{session_cookie}={data}; path={path}; {expires}{security_flags}".format(
  70. session_cookie=self.session_cookie,
  71. data="null",
  72. path=self.path,
  73. expires="expires=Thu, 01 Jan 1970 00:00:00 GMT; ",
  74. security_flags=self.security_flags,
  75. )
  76. headers.append("Set-Cookie", header_value)
  77. await send(message)
  78. await self.app(scope, receive, send_wrapper)
  79. class Session(dict[str, typing.Any]):
  80. accessed: bool = False
  81. modified: bool = False
  82. def mark_accessed(self) -> None:
  83. self.accessed = True
  84. def mark_modified(self) -> None:
  85. self.accessed = True
  86. self.modified = True
  87. def __setitem__(self, key: str, value: typing.Any) -> None:
  88. self.mark_modified()
  89. super().__setitem__(key, value)
  90. def __delitem__(self, key: str) -> None:
  91. self.mark_modified()
  92. super().__delitem__(key)
  93. def clear(self) -> None:
  94. self.mark_modified()
  95. super().clear()
  96. def pop(self, key: str, *args: typing.Any) -> typing.Any:
  97. self.modified = self.modified or key in self
  98. return super().pop(key, *args)
  99. def setdefault(self, key: str, default: typing.Any = None) -> typing.Any:
  100. if key not in self:
  101. self.mark_modified()
  102. return super().setdefault(key, default)
  103. def update(self, *args: typing.Any, **kwargs: typing.Any) -> None:
  104. self.mark_modified()
  105. super().update(*args, **kwargs)