# 定义一个用户表和一个保存用户Token的表

class UserInfo(models.Model):

    username = models.CharField(max_length=16)

    password = models.CharField(max_length=32)

    type = models.SmallIntegerField(

        choices=((1, '普通用户'), (2, 'VIP用户')),

        default=1

    )

class Token(models.Model):

    user = models.OneToOneField(to='UserInfo')

    token_code = models.CharField(max_length=128)

# 视图主要处理用户名、密码是否正确，用户每一次请求都要带着专有token来！

import hashlib, time

from rest_framework.response import Response

from rest_framework.views import APIView

def get_random_token(username):

    """

    根据用户名和时间戳生成随机token

    :param username:

    :return:

    """

    timestamp = str(time.time())

    m = hashlib.md5(bytes(username, encoding="utf8"))

    m.update(bytes(timestamp, encoding="utf8"))

    return m.hexdigest()

class LoginView(APIView):

    """

    校验用户名密码是否正确从而生成token的视图

    """

    def post(self, request):

        res = {"code": 0}

        print(request.data)

        username = request.data.get("username")

        password = request.data.get("password")

        user = models.UserInfo.objects.filter(username=username, password=password).first()

        if user:

            # 如果用户名密码正确

            token = get_random_token(username)

            models.Token.objects.update_or_create(defaults={"token_code": token}, user=user)

            res["token"] = token

        else:

            res["code"] = 1

            res["error"] = "用户名或密码错误"

        return Response(res)

# 这一步是要对着源码才能写出来

from rest_framework.authentication import BaseAuthentication

from rest_framework.exceptions import AuthenticationFailed

class MyAuth(BaseAuthentication):

    def authenticate(self, request):

        if request.method in ["POST", "PUT", "DELETE"]:

            request_token = request.data.get("token", None)

            if not request_token:

                raise AuthenticationFailed('缺少token')

            token_obj = models.Token.objects.filter(token_code=request_token).first()

            if not token_obj:

                raise AuthenticationFailed('无效的token')

            return token_obj.user.username, None

        else:

            return None, None

# 在settings.py中配置

REST_FRAMEWORK = {

    "DEFAULT_AUTHENTICATION_CLASSES": ["app01.utils.MyAuth", ]

}

# 自定义权限类

from rest_framework.permissions import BasePermission

class MyPermission(BasePermission):

    message = 'VIP用户才能访问'

    def has_permission(self, request, view):

        """

        自定义权限只有VIP用户才能访问

        """

        # 因为在进行权限判断之前已经做了认证判断，所以这里可以直接拿到request.user

        if request.user and request.user.type == 2:  # 如果是VIP用户

            return True

        else:

            return False

class CommentViewSet(ModelViewSet):

    queryset = models.Comment.objects.all()

    serializer_class = app01_serializers.CommentSerializer

    authentication_classes = [MyAuth, ]

    permission_classes = [MyPermission, ]

REST_FRAMEWORK = {

    "DEFAULT_AUTHENTICATION_CLASSES": ["app01.utils.MyAuth", ],

    "DEFAULT_PERMISSION_CLASSES": ["app01.utils.MyPermission", ]

}

VISIT_RECORD = {}

# 自定义限制

class MyThrottle(object):

    def __init__(self):

        self.history = None

    def allow_request(self, request, view):

        """

        自定义频率限制60秒内只能访问三次

        """

        # 获取用户IP

        ip = request.META.get("REMOTE_ADDR")

        timestamp = time.time()

        if ip not in VISIT_RECORD:

            VISIT_RECORD[ip] = [timestamp, ]

            return True

        history = VISIT_RECORD[ip]

        self.history = history

        history.insert(0, timestamp)

        while history and history[-1] < timestamp - 60:

            history.pop()

        if len(history) > 3:

            return False

        else:

            return True

    def wait(self):

        """

        限制时间还剩多少

        """

        timestamp = time.time()

        return 60 - (timestamp - self.history[-1])

class CommentViewSet(ModelViewSet):

    queryset = models.Comment.objects.all()

    serializer_class = app01_serializers.CommentSerializer

    throttle_classes = [MyThrottle, ]