FitTrackee/fittrackee/oauth2/resource_protector.py

from functools import wraps
from typing import Any, Callable, List, Union

from authlib.integrations.flask_oauth2 import ResourceProtector
from authlib.oauth2 import OAuth2Error
from authlib.oauth2.rfc6749.errors import MissingAuthorizationError
from flask import current_app, request
from werkzeug.exceptions import RequestEntityTooLarge

from fittrackee.responses import (
    ForbiddenErrorResponse,
    PayloadTooLargeErrorResponse,
    UnauthorizedErrorResponse,
)
from fittrackee.users.models import User


class CustomResourceProtector(ResourceProtector):
    def __call__(
        self,
        scopes: Union[str, List, None] = None,
        as_admin: bool = False,
    ) -> Callable:
        def wrapper(f: Callable) -> Callable:
            @wraps(f)
            def decorated(*args: Any, **kwargs: Any) -> Callable:
                auth_user = None
                auth_header = request.headers.get('Authorization')
                if not auth_header:
                    return UnauthorizedErrorResponse(
                        'provide a valid auth token'
                    )

                # First-party application (Fittrackee front-end)
                # in this case, scopes will be ignored
                auth_token = auth_header.split(' ')[1]
                resp = User.decode_auth_token(auth_token)
                if isinstance(resp, int):
                    auth_user = User.query.filter_by(id=resp).first()

                # Third-party applications
                if not auth_user:
                    current_token = None
                    try:
                        current_token = self.acquire_token(scopes)
                    except MissingAuthorizationError as error:
                        self.raise_error_response(error)
                    except OAuth2Error as error:
                        self.raise_error_response(error)
                    except RequestEntityTooLarge:
                        file_type = ''
                        if request.endpoint in [
                            'auth.edit_picture',
                            'workouts.post_workout',
                        ]:
                            file_type = (
                                'picture'
                                if request.endpoint == 'auth.edit_picture'
                                else 'workout'
                            )
                        return PayloadTooLargeErrorResponse(
                            file_type=file_type,
                            file_size=request.content_length,
                            max_size=current_app.config['MAX_CONTENT_LENGTH'],
                        )
                    auth_user = (
                        None if current_token is None else current_token.user
                    )

                if not auth_user or not auth_user.is_active:
                    return UnauthorizedErrorResponse(
                        'provide a valid auth token'
                    )
                if as_admin and not auth_user.admin:
                    return ForbiddenErrorResponse()
                return f(auth_user, *args, **kwargs)

            return decorated

        return wrapper