from flask import abort
from flask_login import current_user
from flask_socketio import disconnect
from functools import wraps
from .models import Permission


def admin_required(f):
    @wraps(f)
    def wrapped(*args, **kwargs):
        if not current_user.can(Permission.ADMIN):
            abort(403)
        return f(*args, **kwargs)
    return wrapped


def socketio_login_required(f):
    @wraps(f)
    def wrapped(*args, **kwargs):
        if not current_user.is_authenticated:
            disconnect()
        else:
            return f(*args, **kwargs)
    return wrapped


def socketio_admin_required(f):
    @wraps(f)
    def wrapped(*args, **kwargs):
        if not current_user.can(Permission.ADMIN):
            disconnect()
        else:
            return f(*args, **kwargs)
    return wrapped