st-cloud/account/decorators.py

from functools import wraps

from .models import User
from utils.http import make_json_response


def user_passes_test(test_func, error):
    def decorator(view_func):
        @wraps(view_func)
        def _wrapped_view(request, *args, **kwargs):
            if test_func(request):
                return view_func(request, *args, **kwargs)
            return make_json_response(code=401, error=error)
        return _wrapped_view
    return decorator


def login_required(function=None, error='未登录'):
    """
    Decorator for views that checks that the _user is logged in, redirecting
    to the log-in page if necessary.
    """
    def is_login(request):
        if request.method != 'POST':
            return False
        username = request.POST.get('username', '')
        token = request.POST.get('token', '')
        print(f'username={username} token={token}')
        try:
            user = User.objects.get(username=username)
        except:
            print('用户不存在')
            return False
        if not user.check_token(token) or not user.tokens.filter(token=token):
            print('token无效')
            return False
        print('已登录')
        if hasattr(request, 'user'):
            request.user = user
        return True

    actual_decorator = user_passes_test(is_login, error)
    if function:
        return actual_decorator(function)
    return actual_decorator