authentication.py 5.2 KB

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