"""Middleware to handle authentication.""" import base64 import binascii from enum import Enum 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 corvus.api.model import APIMessage from corvus.service import ( authentication_service, user_service, user_token_service ) from corvus.service.role_service import ROLES, Role class Auth(Enum): """Authentication scheme definitions.""" TOKEN = 'TOKEN' BASIC = 'BASIC' NONE = 'NONE' 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: Token header for the username and token. :param header_value: :return: """ if not header_value: return None auth_info = wsgi_to_bytes(header_value) 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('X-Auth-Token', None)) if token and authenticate_with_token(token.username, token.password): return func(*args, **kwargs) return authentication_failed('Token') return decorate def require_role(required_role: Role) -> Callable: """ Decorate require a user role. :param required_role: :return: """ 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 def require(required_auth: Auth, required_role: Role) -> Callable: """ Decorate require Auth and Role :param required_auth: :param required_role: :return: """ def require_decorator(func: Callable) -> Callable: @wraps(func) def decorate(*args: list, **kwargs: dict) -> Any: f = require_role(required_role)(func) if required_auth == Auth.BASIC: f = require_basic_auth(f) elif required_auth == Auth.TOKEN: f = require_token_auth(f) else: Response( response=APIMessage( message="Unexpected Server Error", success=False ), status=500) return f(*args, **kwargs) return decorate return require_decorator