"""Middleware to handle authentication.""" import base64 import binascii from functools import wraps from typing import Optional, Callable, Any from flask import request, Response, g, json 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 ) from atheneum.service.role_service import ROLES, Role 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 authorization_failed(required_role: str) -> Response: """Return a correct response for failed authorization.""" return Response( status=401, response=json.dumps({ 'message': '{} role not present'.format(required_role) }), content_type='application/json' ) def parse_token_header( header_value: str) -> Optional[Authorization]: """ Parse the Authorization: Bearer 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'bearer': try: username, token = base64.b64decode(auth_info).split(b':', 1) except binascii.Error: return None return Authorization('bearer', {'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 def require_role(required_role: Role) -> Callable: """Decorate require user role.""" def required_role_decorator(func: Callable) -> Callable: """Decorate the function.""" @wraps(func) def decorate(*args: list, **kwargs: dict) -> Any: """Require a user role.""" if g.user.role in ROLES.find_roles_in_hierarchy(required_role): return func(*args, **kwargs) return authorization_failed(required_role.value) return decorate return required_role_decorator