"""Middleware to handle authentication.""" import base64 from functools import wraps from typing import Optional, Callable, Any import binascii from flask import request, Response, g from werkzeug.datastructures import Authorization from werkzeug.http import bytes_to_wsgi, wsgi_to_bytes from atheneum.service import ( authentication_service, user_service, user_token_service ) def authenticate_with_password(name: str, password: str) -> bool: """ Authenticate a username and a password. :param name: :param password: :return: """ user = user_service.find_by_name(name) if user is not None \ and authentication_service.is_valid_password(user, password): g.user = user return True return False def authenticate_with_token(name: str, token: str) -> bool: """ Authenticate a username and a token. :param name: :param token: :return: """ user = user_service.find_by_name(name) if user is not None: user_token = user_token_service.find_by_user_and_token(user, token) if user is not None \ and authentication_service.is_valid_token(user_token): g.user = user g.user_token = user_token return True return False def authentication_failed(auth_type: str) -> Response: """ Return a correct response for failed authentication. :param auth_type: :return: """ return Response( status=401, headers={ 'WWW-Authenticate': '%s realm="Login Required"' % auth_type }) def parse_token_header( header_value: str) -> Optional[Authorization]: """ Parse the Authorization: Token header for the username and token. :param header_value: :return: """ if not header_value: return None value = wsgi_to_bytes(header_value) try: auth_type, auth_info = value.split(None, 1) auth_type = auth_type.lower() except ValueError: return None if auth_type == b'token': try: username, token = base64.b64decode(auth_info).split(b':', 1) except binascii.Error: return None return Authorization('token', {'username': bytes_to_wsgi(username), 'password': bytes_to_wsgi(token)}) return None def require_basic_auth(func: Callable) -> Callable: """ Decorate require inline basic auth. :param func: :return: """ @wraps(func) def decorate(*args: list, **kwargs: dict) -> Any: """ Authenticate with a password. :param args: :param kwargs: :return: """ auth = request.authorization if auth and authenticate_with_password(auth.username, auth.password): return func(*args, **kwargs) return authentication_failed('Basic') return decorate def require_token_auth(func: Callable) -> Callable: """ Decorate require inline token auth. :param func: :return: """ @wraps(func) def decorate(*args: list, **kwargs: dict) -> Any: """ Authenticate with a token. :param args: :param kwargs: :return: """ token = parse_token_header( request.headers.get('Authorization', None)) if token and authenticate_with_token(token.username, token.password): return func(*args, **kwargs) return authentication_failed('Bearer') return decorate