32 lines
1.2 KiB
Python
32 lines
1.2 KiB
Python
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.')
|
|
|