nopaque/app/decorators.py

from flask import abort, request
from flask_login import current_user
from functools import wraps
from typing import List, Union
from werkzeug.exceptions import NotAcceptable
from app.models import Permission


def permission_required(permission):
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            if not current_user.can(permission):
                abort(403)
            return f(*args, **kwargs)
        return decorated_function
    return decorator


def admin_required(f):
    return permission_required(Permission.ADMINISTRATE)(f)


def socketio_login_required(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        if current_user.is_authenticated:
            return f(*args, **kwargs)
        return {'code': 401, 'body': 'Unauthorized'}
    return wrapper


def socketio_permission_required(permission):
    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            if not current_user.can(permission):
                return {'code': 403, 'body': 'Forbidden'}
            return f(*args, **kwargs)
        return wrapper
    return decorator


def socketio_admin_required(f):
    return socketio_permission_required(Permission.ADMINISTRATE)(f)


def content_negotiation(
    produces: Union[str, List[str], None] = None,
    consumes: Union[str, List[str], None] = None
):
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            provided = request.mimetype
            if consumes is None:
                consumeables = None
            elif isinstance(consumes, str):
                consumeables = {consumes}
            elif isinstance(consumes, list) and all(isinstance(x, str) for x in consumes):
                consumeables = {*consumes}
            else:
                raise TypeError()
            accepted = {*request.accept_mimetypes.values()}
            if produces is None:
                produceables = None
            elif isinstance(produces, str):
                produceables = {produces}
            elif isinstance(produces, list) and all(isinstance(x, str) for x in produces):
                produceables = {*produces}
            else:
                raise TypeError()
            if produceables is not None and len(produceables & accepted) == 0:
                raise NotAcceptable()
            if consumeables is not None and provided not in consumeables:
                raise NotAcceptable()
            return f(*args, **kwargs)
        return decorated_function
    return decorator