Corvus/server/atheneum/middleware/authentication_middleware.py

import base64
from functools import wraps
from typing import Optional, Callable, Any

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:
    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:
    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 Response(
        status=401,
        headers={
            'WWW-Authenticate': '%s realm="Login Required"' % auth_type
        })


def parse_token_authorization_header(
        header_value: str) -> Optional[Authorization]:
    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 Exception:
            return None
        return Authorization('token', {'username': bytes_to_wsgi(username),
                                       'password': bytes_to_wsgi(token)})
    return None


def require_basic_auth(func: Callable) -> Callable:
    @wraps(func)
    def decorate(*args: list, **kwargs: dict) -> Any:
        auth = request.authorization
        if auth and authenticate_with_password(auth.username, auth.password):
            return func(*args, **kwargs)
        else:
            return authentication_failed('Basic')

    return decorate


def require_token_auth(func: Callable) -> Callable:
    @wraps(func)
    def decorate(*args: list, **kwargs: dict) -> Any:
        token = parse_token_authorization_header(
            request.headers.get('Authorization', None))
        if token and authenticate_with_token(token.username, token.password):
            return func(*args, **kwargs)
        else:
            return authentication_failed('Bearer')

    return decorate