Atheneum/server/atheneum/middleware/authentication_middleware.py

"""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: 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('Token')

    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