"""
A middleware for injecting a bearer token to the Authorizatioin header
"""
import collections.abc
import typing
import urllib.parse
import secure_cookie.cookie # type: ignore
import werkzeug
[docs]class TokenInjectionMiddleware(object):
"""
injects a bearer token to the Authorizatioin header
looks for a token in a secure cookie or the query parameters and
if found, sets the Authorization header before further processing
the request
"""
[docs] def __init__(self, app: typing.Callable, key: str, secret_key: str):
self.app = app
self.key = key
self.secret_key = secret_key
def __call__(
self,
environ: dict[str, typing.Any],
start_response: typing.Callable) -> collections.abc.Iterable:
request = werkzeug.wrappers.Request(environ)
cookie = secure_cookie.cookie.SecureCookie.load_cookie(
request=request, key=self.key, secret_key=self.secret_key)
if cookie_token := cookie.get('token'):
environ['HTTP_AUTHORIZATION'] = f"Bearer {cookie_token['access_token']}"
params = urllib.parse.parse_qs(environ['QUERY_STRING'])
if query_token := params.get('access_token'):
environ['HTTP_AUTHORIZATION'] = f"Bearer {query_token[0]}"
return self.app(environ, start_response)
[docs] def save_token(self, token: typing.Any, response: werkzeug.wrappers.Response) -> None:
"""
saves the provided token in a secure cookie for future requests
"""
if token:
cookie = secure_cookie.cookie.SecureCookie(
data=dict(token=token),
secret_key=self.secret_key,
)
cookie.save_cookie(response, key=self.key, force=True)
else:
cookie = secure_cookie.cookie.SecureCookie(secret_key=self.secret_key)
cookie.save_cookie(response, key=self.key, force=True, expires=0)