from rest_framework.authentication import TokenAuthentication
from rest_framework.exceptions import AuthenticationFailed
from .models import CustomToken
from datetime import timedelta
from django.utils import timezone
from functools import wraps

class CustomTokenAuthentication(TokenAuthentication):
    def authenticate(self, request):
        # 获取请求中的 token
        token_key = request.headers.get('token')
        if token_key:
            # 移除前缀 (如果有 "Token " 前缀的话)
            if token_key.startswith('Token '):
                token_key = token_key[6:]
            # print(token_key)
            try:
                # 获取自定义的 Token 对象
                token = CustomToken.objects.get(key=token_key)
                
                # 检查 token 是否过期
                if token.expiration_date < timezone.now():
                    raise AuthenticationFailed('Token has expired.')

                return (token.user, token)  # 返回用户和 token

            except CustomToken.DoesNotExist:
                raise AuthenticationFailed('Invalid token.')

        raise AuthenticationFailed('Authorization header missing or invalid.')