st-cloud/account/decorators.py

from functools import wraps
from urllib.parse import urlparse

from django.conf import settings
from django.shortcuts import resolve_url
from django.http import JsonResponse

from .models import User, LoginToken


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 JsonResponse({'code': 401, 'error': error}, status=401)
        return _wrapped_view
    return decorator


def login_required(function=None, error='error'):
    """
    Decorator for views that checks that the _user is logged in, redirecting
    to the log-in page if necessary.
    """
    def is_login(request):
        username = request.data.get('username', '')
        token = request.data.get('token', '')
        try:
            user = User.objects.get(username='username')
            if user.check_token(token):
                user.tokens.get(token=token)
                return True
        except:
            return False
        return False

    actual_decorator = user_passes_test(is_login, '请登录')
    if function:
        return actual_decorator(function)
    return actual_decorator