import logging
|
from datetime import datetime
|
|
from core.permissions import ViewClassPermission, all_permissions
|
from django.utils.decorators import method_decorator
|
from drf_spectacular.utils import extend_schema
|
from jwt_auth.auth import TokenAuthenticationPhaseout
|
from jwt_auth.models import LSAPIToken, TruncatedLSAPIToken
|
from jwt_auth.serializers import (
|
JWTSettingsSerializer,
|
LSAPITokenCreateSerializer,
|
LSAPITokenListSerializer,
|
TokenRefreshResponseSerializer,
|
TokenRotateResponseSerializer,
|
)
|
from rest_framework import generics, status
|
from rest_framework.authentication import SessionAuthentication
|
from rest_framework.exceptions import APIException
|
from rest_framework.generics import CreateAPIView
|
from rest_framework.response import Response
|
from rest_framework.settings import api_settings
|
from rest_framework_simplejwt.authentication import JWTAuthentication
|
from rest_framework_simplejwt.exceptions import TokenBackendError, TokenError
|
from rest_framework_simplejwt.token_blacklist.models import BlacklistedToken, OutstandingToken
|
from rest_framework_simplejwt.views import TokenRefreshView, TokenViewBase
|
|
logger = logging.getLogger(__name__)
|
|
|
class TokenExistsError(APIException):
|
status_code = status.HTTP_409_CONFLICT
|
default_detail = 'You already have a valid token. Please revoke it before creating a new one.'
|
default_code = 'token_exists'
|
|
|
@method_decorator(
|
name='get',
|
decorator=extend_schema(
|
tags=['JWT'],
|
summary='Retrieve JWT Settings',
|
description='Retrieve JWT settings for the currently active organization.',
|
extensions={
|
'x-fern-sdk-group-name': 'jwt_settings',
|
'x-fern-sdk-method-name': 'get',
|
'x-fern-audiences': ['public'],
|
},
|
),
|
)
|
@method_decorator(
|
name='post',
|
decorator=extend_schema(
|
tags=['JWT'],
|
summary='Update JWT Settings',
|
description='Update JWT settings for the currently active organization.',
|
extensions={
|
'x-fern-sdk-group-name': 'jwt_settings',
|
'x-fern-sdk-method-name': 'update',
|
'x-fern-audiences': ['public'],
|
},
|
),
|
)
|
class JWTSettingsAPI(CreateAPIView):
|
serializer_class = JWTSettingsSerializer
|
permission_required = ViewClassPermission(
|
GET=all_permissions.organizations_view,
|
POST=all_permissions.organizations_change,
|
)
|
|
def get_object(self):
|
jwt = self.request.user.active_organization.jwt
|
self.check_object_permissions(self.request, jwt)
|
return jwt
|
|
def get(self, request, *args, **kwargs):
|
jwt_settings = self.get_object()
|
return Response(self.get_serializer(jwt_settings).data)
|
|
def post(self, request, *args, **kwargs):
|
jwt_settings = self.get_object()
|
serializer = self.get_serializer(data=request.data, instance=jwt_settings)
|
serializer.is_valid(raise_exception=True)
|
serializer.save()
|
return Response(serializer.data)
|
|
|
class DecoratedTokenRefreshView(TokenRefreshView):
|
@extend_schema(
|
tags=['JWT'],
|
summary='Refresh JWT token',
|
description='Get a new access token, using a refresh token.',
|
responses={
|
status.HTTP_200_OK: TokenRefreshResponseSerializer,
|
},
|
extensions={
|
'x-fern-sdk-group-name': 'tokens',
|
'x-fern-sdk-method-name': 'refresh',
|
'x-fern-audiences': ['public'],
|
},
|
)
|
def post(self, request, *args, **kwargs):
|
return super().post(request, *args, **kwargs)
|
|
|
@method_decorator(
|
name='get',
|
decorator=extend_schema(
|
tags=['JWT'],
|
summary='List API tokens',
|
description='List all API tokens for the current user.',
|
responses={
|
status.HTTP_200_OK: LSAPITokenListSerializer,
|
},
|
extensions={
|
'x-fern-sdk-group-name': 'tokens',
|
'x-fern-sdk-method-name': 'list',
|
'x-fern-audiences': ['public'],
|
},
|
),
|
)
|
@method_decorator(
|
name='post',
|
decorator=extend_schema(
|
tags=['JWT'],
|
summary='Create API token',
|
description='Create a new API token for the current user.',
|
responses={
|
status.HTTP_201_CREATED: LSAPITokenCreateSerializer,
|
},
|
extensions={
|
'x-fern-sdk-group-name': 'tokens',
|
'x-fern-sdk-method-name': 'create',
|
'x-fern-audiences': ['public'],
|
},
|
),
|
)
|
class LSAPITokenView(generics.ListCreateAPIView):
|
permission_required = all_permissions.users_token_any
|
token_class = LSAPIToken
|
|
def get_queryset(self):
|
"""Returns all non-expired non-blacklisted tokens for the current user.
|
|
The `list` method handles filtering for refresh tokens (as opposed to access tokens),
|
since simple-jwt makes it hard to do this at the DB level."""
|
# Notably, if the list of non-expired blacklisted tokens ever gets too long
|
# (e.g. users from orgs who have not set a token expiration for their org
|
# revoke enough tokens for this to blow up), this will become inefficient.
|
# Would be ideal to just add a "blacklisted" attr to our own subclass of
|
# OutstandingToken so we can check at that level, or just clean up
|
# OutstandingTokens that have been blacklisted every so often.
|
current_blacklisted_tokens = BlacklistedToken.objects.filter(token__expires_at__gt=datetime.now()).values_list(
|
'token_id', flat=True
|
)
|
return OutstandingToken.objects.filter(user_id=self.request.user.id, expires_at__gt=datetime.now()).exclude(
|
id__in=current_blacklisted_tokens
|
)
|
|
def list(self, request, *args, **kwargs):
|
all_tokens = self.get_queryset()
|
|
def _maybe_get_token(token: OutstandingToken):
|
try:
|
return TruncatedLSAPIToken(str(token.token))
|
except (TokenError, TokenBackendError) as e: # expired/invalid token
|
logger.debug('JWT API token validation failed: %s', e)
|
return None
|
|
# Annoyingly, token_type not stored directly so we have to filter it here.
|
# Shouldn't be many unexpired tokens to iterate through.
|
token_objects = list(filter(None, [_maybe_get_token(token) for token in all_tokens]))
|
refresh_tokens = [tok for tok in token_objects if tok['token_type'] == 'refresh']
|
|
serializer = self.get_serializer(refresh_tokens, many=True)
|
data = serializer.data
|
return Response(data)
|
|
def get_serializer_class(self):
|
if self.request.method == 'POST':
|
return LSAPITokenCreateSerializer
|
return LSAPITokenListSerializer
|
|
def perform_create(self, serializer):
|
# Check for existing valid tokens
|
existing_tokens = self.get_queryset()
|
if existing_tokens.exists():
|
raise TokenExistsError()
|
|
token = self.token_class.for_user(self.request.user)
|
serializer.instance = token
|
|
|
class LSTokenBlacklistView(TokenViewBase):
|
_serializer_class = 'jwt_auth.serializers.LSAPITokenBlacklistSerializer'
|
|
@extend_schema(
|
tags=['JWT'],
|
summary='Blacklist a JWT refresh token',
|
description='Adds a JWT refresh token to the blacklist, preventing it from being used to obtain new access tokens.',
|
responses={
|
status.HTTP_204_NO_CONTENT: 'Token was successfully blacklisted',
|
status.HTTP_404_NOT_FOUND: 'Token is already blacklisted',
|
},
|
extensions={
|
'x-fern-sdk-group-name': 'tokens',
|
'x-fern-sdk-method-name': 'blacklist',
|
'x-fern-audiences': ['public'],
|
},
|
)
|
def post(self, request, *args, **kwargs):
|
serializer = self.get_serializer(data=request.data)
|
try:
|
# Notably, simple jwt's serializer (which we inherit from) calls
|
# .blacklist() on the token under the hood
|
serializer.is_valid(raise_exception=True)
|
except TokenError as e:
|
logger.error('Token error occurred while trying to blacklist a token: %s', str(e), exc_info=True)
|
return Response({'detail': 'Token is invalid or already blacklisted.'}, status=status.HTTP_404_NOT_FOUND)
|
|
return Response(status=status.HTTP_204_NO_CONTENT)
|
|
|
class LSAPITokenRotateView(TokenViewBase):
|
# Have to explicitly set authentication_classes here, due to how auth works in our middleware, request.user is not set
|
# properly before executing the view.
|
authentication_classes = [JWTAuthentication, TokenAuthenticationPhaseout, SessionAuthentication]
|
permission_classes = api_settings.DEFAULT_PERMISSION_CLASSES
|
permission_required = all_permissions.users_token_any
|
_serializer_class = 'jwt_auth.serializers.LSAPITokenRotateSerializer'
|
token_class = LSAPIToken
|
|
@extend_schema(
|
tags=['JWT'],
|
summary='Rotate JWT refresh token',
|
description='Creates a new JWT refresh token and blacklists the current one.',
|
responses={
|
status.HTTP_200_OK: TokenRotateResponseSerializer,
|
status.HTTP_400_BAD_REQUEST: 'Invalid token or token already blacklisted',
|
},
|
extensions={
|
'x-fern-sdk-group-name': 'tokens',
|
'x-fern-sdk-method-name': 'rotate',
|
'x-fern-audiences': ['public'],
|
},
|
)
|
def post(self, request, *args, **kwargs):
|
serializer = self.get_serializer(data=request.data)
|
serializer.is_valid(raise_exception=True)
|
current_token = serializer.validated_data['refresh']
|
|
# Blacklist the current token
|
try:
|
current_token.blacklist()
|
except TokenError:
|
return Response({'detail': 'Token is invalid or already blacklisted.'}, status=status.HTTP_400_BAD_REQUEST)
|
|
# Create a new token for the user
|
new_token = self.create_token(request.user)
|
return Response({'refresh': new_token.get_full_jwt()}, status=status.HTTP_200_OK)
|
|
def create_token(self, user):
|
"""Create a new token for the user. Can be overridden by child classes to use different token classes."""
|
return self.token_class.for_user(user)
|
