authentication.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. from __future__ import annotations
  2. import functools
  3. import inspect
  4. from collections.abc import Callable, Sequence
  5. from typing import Any, ParamSpec
  6. from urllib.parse import urlencode
  7. from starlette._utils import is_async_callable
  8. from starlette.exceptions import HTTPException
  9. from starlette.requests import HTTPConnection, Request
  10. from starlette.responses import RedirectResponse
  11. from starlette.websockets import WebSocket
  12. _P = ParamSpec("_P")
  13. def has_required_scope(conn: HTTPConnection, scopes: Sequence[str]) -> bool:
  14. for scope in scopes:
  15. if scope not in conn.auth.scopes:
  16. return False
  17. return True
  18. def requires(
  19. scopes: str | Sequence[str],
  20. status_code: int = 403,
  21. redirect: str | None = None,
  22. ) -> Callable[[Callable[_P, Any]], Callable[_P, Any]]:
  23. scopes_list = [scopes] if isinstance(scopes, str) else list(scopes)
  24. def decorator(
  25. func: Callable[_P, Any],
  26. ) -> Callable[_P, Any]:
  27. sig = inspect.signature(func)
  28. for idx, parameter in enumerate(sig.parameters.values()):
  29. if parameter.name == "request" or parameter.name == "websocket":
  30. type_ = parameter.name
  31. break
  32. else:
  33. raise Exception(f'No "request" or "websocket" argument on function "{func}"')
  34. if type_ == "websocket":
  35. # Handle websocket functions. (Always async)
  36. @functools.wraps(func)
  37. async def websocket_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
  38. websocket = kwargs.get("websocket", args[idx] if idx < len(args) else None)
  39. assert isinstance(websocket, WebSocket)
  40. if not has_required_scope(websocket, scopes_list):
  41. await websocket.close()
  42. else:
  43. await func(*args, **kwargs)
  44. return websocket_wrapper
  45. elif is_async_callable(func):
  46. # Handle async request/response functions.
  47. @functools.wraps(func)
  48. async def async_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> Any:
  49. request = kwargs.get("request", args[idx] if idx < len(args) else None)
  50. assert isinstance(request, Request)
  51. if not has_required_scope(request, scopes_list):
  52. if redirect is not None:
  53. orig_request_qparam = urlencode({"next": str(request.url)})
  54. next_url = f"{request.url_for(redirect)}?{orig_request_qparam}"
  55. return RedirectResponse(url=next_url, status_code=303)
  56. raise HTTPException(status_code=status_code)
  57. return await func(*args, **kwargs)
  58. return async_wrapper
  59. else:
  60. # Handle sync request/response functions.
  61. @functools.wraps(func)
  62. def sync_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> Any:
  63. request = kwargs.get("request", args[idx] if idx < len(args) else None)
  64. assert isinstance(request, Request)
  65. if not has_required_scope(request, scopes_list):
  66. if redirect is not None:
  67. orig_request_qparam = urlencode({"next": str(request.url)})
  68. next_url = f"{request.url_for(redirect)}?{orig_request_qparam}"
  69. return RedirectResponse(url=next_url, status_code=303)
  70. raise HTTPException(status_code=status_code)
  71. return func(*args, **kwargs)
  72. return sync_wrapper
  73. return decorator
  74. class AuthenticationError(Exception):
  75. pass
  76. class AuthenticationBackend:
  77. async def authenticate(self, conn: HTTPConnection) -> tuple[AuthCredentials, BaseUser] | None:
  78. raise NotImplementedError() # pragma: no cover
  79. class AuthCredentials:
  80. def __init__(self, scopes: Sequence[str] | None = None):
  81. self.scopes = [] if scopes is None else list(scopes)
  82. class BaseUser:
  83. @property
  84. def is_authenticated(self) -> bool:
  85. raise NotImplementedError() # pragma: no cover
  86. @property
  87. def display_name(self) -> str:
  88. raise NotImplementedError() # pragma: no cover
  89. @property
  90. def identity(self) -> str:
  91. raise NotImplementedError() # pragma: no cover
  92. class SimpleUser(BaseUser):
  93. def __init__(self, username: str) -> None:
  94. self.username = username
  95. @property
  96. def is_authenticated(self) -> bool:
  97. return True
  98. @property
  99. def display_name(self) -> str:
  100. return self.username
  101. class UnauthenticatedUser(BaseUser):
  102. @property
  103. def is_authenticated(self) -> bool:
  104. return False
  105. @property
  106. def display_name(self) -> str:
  107. return ""