from __future__ import annotations from django.utils import timezone from rest_framework.authentication import BaseAuthentication, get_authorization_header from rest_framework.exceptions import AuthenticationFailed, PermissionDenied from apps.api.errors import api_error from apps.users.models import ApiKey class ApiKeyAuthentication(BaseAuthentication): keyword = "Bearer" def authenticate(self, request): raw_header = get_authorization_header(request) if not raw_header: return None try: header = raw_header.decode("utf-8") except UnicodeError as exc: raise self.authentication_failed() from exc parts = header.split() if len(parts) != 2 or parts[0].lower() != self.keyword.lower(): raise self.authentication_failed() raw_key = parts[1] if not raw_key: raise self.authentication_failed() key_hash = ApiKey.hash_key(raw_key) try: api_key = ApiKey.objects.select_related("user").get(key_hash=key_hash) except ApiKey.DoesNotExist as exc: raise self.authentication_failed() from exc if not api_key.matches_key(raw_key): raise self.authentication_failed() if not api_key.is_active_key: raise PermissionDenied( api_error("account_disabled", "账号或 API Key 已禁用") ) if not api_key.user.is_business_active: raise PermissionDenied( api_error("account_disabled", "账号或 API Key 已禁用") ) now = timezone.now() ApiKey.objects.filter(pk=api_key.pk).update(last_used_at=now) api_key.last_used_at = now return api_key.user, api_key def authenticate_header(self, request) -> str: return self.keyword @staticmethod def authentication_failed() -> AuthenticationFailed: return AuthenticationFailed(api_error("unauthorized", "缺失或无效 API Key"))