1  from datetime import datetime
   2  from datetime import timezone
   3  from functools import wraps
   4  from re import split
   5  from typing import Any
   6  from typing import Optional
   7  from typing import Sequence
   8  from typing import Tuple
   9  from typing import Union
  10  
  11  from flask import current_app
  12  from flask import g
  13  from flask import request
  14  from werkzeug.exceptions import BadRequest
  15  
  16  from flask_jwt_extended.config import config
  17  from flask_jwt_extended.exceptions import CSRFError
  18  from flask_jwt_extended.exceptions import FreshTokenRequired
  19  from flask_jwt_extended.exceptions import InvalidHeaderError
  20  from flask_jwt_extended.exceptions import InvalidQueryParamError
  21  from flask_jwt_extended.exceptions import NoAuthorizationError
  22  from flask_jwt_extended.exceptions import UserLookupError
  23  from flask_jwt_extended.internal_utils import custom_verification_for_token
  24  from flask_jwt_extended.internal_utils import has_user_lookup
  25  from flask_jwt_extended.internal_utils import user_lookup
  26  from flask_jwt_extended.internal_utils import verify_token_not_blocklisted
  27  from flask_jwt_extended.internal_utils import verify_token_type
  28  from flask_jwt_extended.utils import decode_token
  29  from flask_jwt_extended.utils import get_unverified_jwt_headers
  30  
  31  LocationType = Union[str, Sequence, None]
  32  
  33  
  34  def _verify_token_is_fresh(jwt_header: dict, jwt_data: dict) -> None:
  35      fresh = jwt_data["fresh"]
  36      if isinstance(fresh, bool):
  37          if not fresh:
  38              raise FreshTokenRequired("Fresh token required", jwt_header, jwt_data)
  39      else:
  40          now = datetime.timestamp(datetime.now(timezone.utc))
  41          if fresh < now:
  42              raise FreshTokenRequired("Fresh token required", jwt_header, jwt_data)
  43  
  44  
  45  def verify_jwt_in_request(
  46      optional: bool = False,
  47      fresh: bool = False,
  48      refresh: bool = False,
  49      locations: Optional[LocationType] = None,
  50      verify_type: bool = True,
  51      skip_revocation_check: bool = False,
  52  ) -> Optional[Tuple[dict, dict]]:
  53      """
  54      Verify that a valid JWT is present in the request, unless ``optional=True`` in
  55      which case no JWT is also considered valid.
  56  
  57      :param optional:
  58          If ``True``, do not raise an error if no JWT is present in the request.
  59          Defaults to ``False``.
  60  
  61      :param fresh:
  62          If ``True``, require a JWT marked as ``fresh`` in order to be verified.
  63          Defaults to ``False``.
  64  
  65      :param refresh:
  66          If ``True``, requires a refresh JWT to access this endpoint. If ``False``,
  67          requires an access JWT to access this endpoint. Defaults to ``False``
  68  
  69      :param locations:
  70          A location or list of locations to look for the JWT in this request, for
  71          example ``'headers'`` or ``['headers', 'cookies']``. Defaults to ``None``
  72          which indicates that JWTs will be looked for in the locations defined by the
  73          ``JWT_TOKEN_LOCATION`` configuration option.
  74  
  75      :param verify_type:
  76          If ``True``, the token type (access or refresh) will be checked according
  77          to the ``refresh`` argument. If ``False``, type will not be checked and both
  78          access and refresh tokens will be accepted.
  79  
  80      :param skip_revocation_check:
  81          If ``True``, revocation status of the token will be *not* checked. If ``False``,
  82          revocation status of the token will be checked.
  83  
  84      :return:
  85          A tuple containing the jwt_header and the jwt_data if a valid JWT is
  86          present in the request. If ``optional=True`` and no JWT is in the request,
  87          ``None`` will be returned instead. Raise an exception if an invalid JWT
  88          is in the request.
  89      """
  90      if request.method in config.exempt_methods:
  91          return None
  92  
  93      try:
  94          jwt_data, jwt_header, jwt_location = _decode_jwt_from_request(
  95              locations,
  96              fresh,
  97              refresh=refresh,
  98              verify_type=verify_type,
  99              skip_revocation_check=skip_revocation_check,
 100          )
 101  
 102      except NoAuthorizationError:
 103          if not optional:
 104              raise
 105          g._jwt_extended_jwt = {}
 106          g._jwt_extended_jwt_header = {}
 107          g._jwt_extended_jwt_user = {"loaded_user": None}
 108          g._jwt_extended_jwt_location = None
 109          return None
 110  
 111      # Save these at the very end so that they are only saved in the requet
 112      # context if the token is valid and all callbacks succeed
 113      g._jwt_extended_jwt_user = _load_user(jwt_header, jwt_data)
 114      g._jwt_extended_jwt_header = jwt_header
 115      g._jwt_extended_jwt = jwt_data
 116      g._jwt_extended_jwt_location = jwt_location
 117  
 118      return jwt_header, jwt_data
 119  
 120  
 121  def jwt_required(
 122      optional: bool = False,
 123      fresh: bool = False,
 124      refresh: bool = False,
 125      locations: Optional[LocationType] = None,
 126      verify_type: bool = True,
 127      skip_revocation_check: bool = False,
 128  ) -> Any:
 129      """
 130      A decorator to protect a Flask endpoint with JSON Web Tokens.
 131  
 132      Any route decorated with this will require a valid JWT to be present in the
 133      request (unless optional=True, in which case no JWT is also valid) before the
 134      endpoint can be called.
 135  
 136      :param optional:
 137          If ``True``, allow the decorated endpoint to be accessed if no JWT is present in
 138          the request. Defaults to ``False``.
 139  
 140      :param fresh:
 141          If ``True``, require a JWT marked with ``fresh`` to be able to access this
 142          endpoint. Defaults to ``False``.
 143  
 144      :param refresh:
 145          If ``True``, requires a refresh JWT to access this endpoint. If ``False``,
 146          requires an access JWT to access this endpoint. Defaults to ``False``.
 147  
 148      :param locations:
 149          A location or list of locations to look for the JWT in this request, for
 150          example ``'headers'`` or ``['headers', 'cookies']``. Defaults to ``None``
 151          which indicates that JWTs will be looked for in the locations defined by the
 152          ``JWT_TOKEN_LOCATION`` configuration option.
 153  
 154      :param verify_type:
 155          If ``True``, the token type (access or refresh) will be checked according
 156          to the ``refresh`` argument. If ``False``, type will not be checked and both
 157          access and refresh tokens will be accepted.
 158  
 159      :param skip_revocation_check:
 160          If ``True``, revocation status of the token will be *not* checked. If ``False``,
 161          revocation status of the token will be checked.
 162      """
 163  
 164      def wrapper(fn):
 165          @wraps(fn)
 166          def decorator(*args, **kwargs):
 167              verify_jwt_in_request(
 168                  optional, fresh, refresh, locations, verify_type, skip_revocation_check
 169              )
 170              return current_app.ensure_sync(fn)(*args, **kwargs)
 171  
 172          return decorator
 173  
 174      return wrapper
 175  
 176  
 177  def _load_user(jwt_header: dict, jwt_data: dict) -> Optional[dict]:
 178      if not has_user_lookup():
 179          return None
 180  
 181      identity = jwt_data[config.identity_claim_key]
 182      user = user_lookup(jwt_header, jwt_data)
 183      if user is None:
 184          error_msg = "user_lookup returned None for {}".format(identity)
 185          raise UserLookupError(error_msg, jwt_header, jwt_data)
 186      return {"loaded_user": user}
 187  
 188  
 189  def _decode_jwt_from_headers() -> Tuple[str, None]:
 190      header_name = config.header_name
 191      header_type = config.header_type
 192  
 193      # Verify we have the auth header
 194      auth_header = request.headers.get(header_name, "").strip().strip(",")
 195      if not auth_header:
 196          raise NoAuthorizationError(f"Missing {header_name} Header")
 197  
 198      # Make sure the header is in a valid format that we are expecting, ie
 199      # <HeaderName>: <HeaderType(optional)> <JWT>.
 200      #
 201      # Also handle the fact that the header that can be comma delimited, ie
 202      # <HeaderName>: <field> <value>, <field> <value>, etc...
 203      if header_type:
 204          field_values = split(r",\s*", auth_header)
 205          jwt_headers = [s for s in field_values if s.split()[0] == header_type]
 206          if len(jwt_headers) != 1:
 207              msg = (
 208                  f"Missing '{header_type}' type in '{header_name}' header. "
 209                  f"Expected '{header_name}: {header_type} <JWT>'"
 210              )
 211              raise NoAuthorizationError(msg)
 212  
 213          parts = jwt_headers[0].split()
 214          if len(parts) != 2:
 215              msg = (
 216                  f"Bad {header_name} header. "
 217                  f"Expected '{header_name}: {header_type} <JWT>'"
 218              )
 219              raise InvalidHeaderError(msg)
 220  
 221          encoded_token = parts[1]
 222      else:
 223          parts = auth_header.split()
 224          if len(parts) != 1:
 225              msg = f"Bad {header_name} header. Expected '{header_name}: <JWT>'"
 226              raise InvalidHeaderError(msg)
 227  
 228          encoded_token = parts[0]
 229  
 230      return encoded_token, None
 231  
 232  
 233  def _decode_jwt_from_cookies(refresh: bool) -> Tuple[str, Optional[str]]:
 234      if refresh:
 235          cookie_key = config.refresh_cookie_name
 236          csrf_header_key = config.refresh_csrf_header_name
 237          csrf_field_key = config.refresh_csrf_field_name
 238      else:
 239          cookie_key = config.access_cookie_name
 240          csrf_header_key = config.access_csrf_header_name
 241          csrf_field_key = config.access_csrf_field_name
 242  
 243      encoded_token = request.cookies.get(cookie_key)
 244      if not encoded_token:
 245          raise NoAuthorizationError('Missing cookie "{}"'.format(cookie_key))
 246  
 247      if config.csrf_protect and request.method in config.csrf_request_methods:
 248          csrf_value = request.headers.get(csrf_header_key, None)
 249          if not csrf_value and config.csrf_check_form:
 250              csrf_value = request.form.get(csrf_field_key, None)
 251          if not csrf_value:
 252              raise CSRFError("Missing CSRF token")
 253      else:
 254          csrf_value = None
 255  
 256      return encoded_token, csrf_value
 257  
 258  
 259  def _decode_jwt_from_query_string() -> Tuple[str, None]:
 260      param_name = config.query_string_name
 261      prefix = config.query_string_value_prefix
 262  
 263      value = request.args.get(param_name)
 264      if not value:
 265          raise NoAuthorizationError(f"Missing '{param_name}' query paramater")
 266  
 267      if not value.startswith(prefix):
 268          raise InvalidQueryParamError(
 269              f"Invalid value for query parameter '{param_name}'. "
 270              f"Expected the value to start with '{prefix}'"
 271          )
 272  
 273      encoded_token = value[len(prefix) :]  # noqa: E203
 274      return encoded_token, None
 275  
 276  
 277  def _decode_jwt_from_json(refresh: bool) -> Tuple[str, None]:
 278      if not request.is_json:
 279          raise NoAuthorizationError("Invalid content-type. Must be application/json.")
 280  
 281      if refresh:
 282          token_key = config.refresh_json_key
 283      else:
 284          token_key = config.json_key
 285  
 286      try:
 287          encoded_token = request.json and request.json.get(token_key, None)
 288          if not encoded_token:
 289              raise BadRequest()
 290      except BadRequest:
 291          raise NoAuthorizationError(
 292              'Missing "{}" key in json data.'.format(token_key)
 293          ) from None
 294  
 295      return encoded_token, None
 296  
 297  
 298  def _decode_jwt_from_request(
 299      locations: LocationType,
 300      fresh: bool,
 301      refresh: bool = False,
 302      verify_type: bool = True,
 303      skip_revocation_check: bool = False,
 304  ) -> Tuple[dict, dict, str]:
 305      # Figure out what locations to look for the JWT in this request
 306      if isinstance(locations, str):
 307          locations = [locations]
 308  
 309      if not locations:
 310          locations = config.token_location
 311  
 312      # Get the decode functions in the order specified by locations.
 313      # Each entry in this list is a tuple (<location>, <encoded-token-function>)
 314      get_encoded_token_functions = []
 315      for location in locations:
 316          if location == "cookies":
 317              get_encoded_token_functions.append(
 318                  (location, lambda: _decode_jwt_from_cookies(refresh))
 319              )
 320          elif location == "query_string":
 321              get_encoded_token_functions.append(
 322                  (location, _decode_jwt_from_query_string)
 323              )
 324          elif location == "headers":
 325              get_encoded_token_functions.append((location, _decode_jwt_from_headers))
 326          elif location == "json":
 327              get_encoded_token_functions.append(
 328                  (location, lambda: _decode_jwt_from_json(refresh))
 329              )
 330          else:
 331              raise RuntimeError(f"'{location}' is not a valid location")
 332  
 333      # Try to find the token from one of these locations. It only needs to exist
 334      # in one place to be valid (not every location).
 335      errors = []
 336      decoded_token = None
 337      for location, get_encoded_token_function in get_encoded_token_functions:
 338          try:
 339              encoded_token, csrf_token = get_encoded_token_function()
 340              decoded_token = decode_token(encoded_token, csrf_token)
 341              jwt_location = location
 342              jwt_header = get_unverified_jwt_headers(encoded_token)
 343              break
 344          except NoAuthorizationError as e:
 345              errors.append(str(e))
 346  
 347      # Do some work to make a helpful and human readable error message if no
 348      # token was found in any of the expected locations.
 349      if not decoded_token:
 350          if len(locations) > 1:
 351              err_msg = "Missing JWT in {start_locs} or {end_locs} ({details})".format(
 352                  start_locs=", ".join(locations[:-1]),
 353                  end_locs=locations[-1],
 354                  details="; ".join(errors),
 355              )
 356              raise NoAuthorizationError(err_msg)
 357          else:
 358              raise NoAuthorizationError(errors[0])
 359  
 360      # Additional verifications provided by this extension
 361      if verify_type:
 362          verify_token_type(decoded_token, refresh)
 363  
 364      if fresh:
 365          _verify_token_is_fresh(jwt_header, decoded_token)
 366  
 367      if not skip_revocation_check:
 368          verify_token_not_blocklisted(jwt_header, decoded_token)
 369  
 370      custom_verification_for_token(jwt_header, decoded_token)
 371  
 372      return decoded_token, jwt_header, jwt_location
 373