Atheneum/server/atheneum/authentication.py

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

from flask import request, Response, session
from werkzeug.datastructures import Authorization
from werkzeug.http import bytes_to_wsgi, wsgi_to_bytes


def authenticate_with_password(username: str, password: str) -> bool:
    session['user'] = None
    return True


def authenticate_with_token(username: str, token: str) -> bool:
    session['user'] = None
    return True


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) -> Optional[Authorization]:
    if not header_value:
        return
    value = wsgi_to_bytes(header_value)
    try:
        auth_type, auth_info = value.split(None, 1)
        auth_type = auth_type.lower()
    except ValueError:
        return
    if auth_type == b'token':
        try:
            username, token = base64.b64decode(auth_info).split(b':', 1)
        except Exception:
            return
        return Authorization('token', {'username': bytes_to_wsgi(username),
                                       'password': bytes_to_wsgi(token)})


def require_basic_auth(func: Callable) -> Callable:
    @wraps(func)
    def decorate(*args, **kwargs):
        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, **kwargs):
        token = parse_token_authorization_header(
            request.headers.get('WWW-Authenticate', None))
        if token and authenticate_with_token(token.username, token.password):
            return func(*args, **kwargs)
        else:
            return authentication_failed('Token')

    return decorate


def generate_token() -> uuid.UUID:
    return uuid.uuid4()