From 7848beedceb5b5e69d651c4a400a510fa83f18ec Mon Sep 17 00:00:00 2001
From: Abhimanyu Saharan <desk.abhimanyu@gmail.com>
Date: Thu, 31 Aug 2023 01:18:18 +0530
Subject: [PATCH] adds additional parameters for token provision api #12870

---
 netbox/users/api/serializers.py | 42 +++++++++++++++++++++++++++++----
 netbox/users/api/views.py       | 34 +++++++++-----------------
 2 files changed, 49 insertions(+), 27 deletions(-)

diff --git a/netbox/users/api/serializers.py b/netbox/users/api/serializers.py
index 1f4bf4ea0..75ab877cf 100644
--- a/netbox/users/api/serializers.py
+++ b/netbox/users/api/serializers.py
@@ -1,11 +1,12 @@
 from django.conf import settings
+from django.contrib.auth import authenticate
 from django.contrib.auth import get_user_model
 from django.contrib.auth.models import Group
 from django.contrib.contenttypes.models import ContentType
 from drf_spectacular.utils import extend_schema_field
 from drf_spectacular.types import OpenApiTypes
 from rest_framework import serializers
-from rest_framework.exceptions import PermissionDenied
+from rest_framework.exceptions import AuthenticationFailed, PermissionDenied
 
 from netbox.api.fields import ContentTypeField, IPNetworkSerializer, SerializedPKRelatedField
 from netbox.api.serializers import ValidatedModelSerializer
@@ -107,9 +108,42 @@ class TokenSerializer(ValidatedModelSerializer):
         return super().validate(data)
 
 
-class TokenProvisionSerializer(serializers.Serializer):
-    username = serializers.CharField()
-    password = serializers.CharField()
+class TokenProvisionSerializer(TokenSerializer):
+    user = NestedUserSerializer(
+        read_only=True
+    )
+    username = serializers.CharField(
+        write_only=True
+    )
+    password = serializers.CharField(
+        write_only=True
+    )
+    last_used = serializers.DateTimeField(
+        read_only=True
+    )
+    key = serializers.CharField(
+        read_only=True
+    )
+
+    class Meta:
+        model = Token
+        fields = (
+            'id', 'url', 'display', 'user', 'created', 'expires', 'last_used', 'key', 'write_enabled', 'description',
+            'allowed_ips', 'username', 'password',
+        )
+
+    def validate(self, data):
+        # Validate the username and password
+        username = data.pop('username')
+        password = data.pop('password')
+        user = authenticate(request=self.context.get('request'), username=username, password=password)
+        if user is None:
+            raise AuthenticationFailed("Invalid username/password")
+
+        # Inject the user into the validated data
+        data['user'] = user
+
+        return data
 
 
 class ObjectPermissionSerializer(ValidatedModelSerializer):
diff --git a/netbox/users/api/views.py b/netbox/users/api/views.py
index 9cf5b1ac5..62a32c71b 100644
--- a/netbox/users/api/views.py
+++ b/netbox/users/api/views.py
@@ -1,3 +1,4 @@
+import logging
 from django.contrib.auth import authenticate
 from django.contrib.auth import get_user_model
 from django.contrib.auth.models import Group
@@ -63,34 +64,21 @@ class TokenProvisionView(APIView):
     @extend_schema(
         request=serializers.TokenProvisionSerializer,
         responses={
-            201: serializers.TokenSerializer,
+            201: serializers.TokenProvisionSerializer,
             401: OpenApiTypes.OBJECT,
         }
     )
     def post(self, request):
-        serializer = serializers.TokenProvisionSerializer(data=request.data)
-        serializer.is_valid()
+        serializer = serializers.TokenProvisionSerializer(data=request.data, context={'request': request})
+        serializer.is_valid(raise_exception=True)
+        self.perform_create(serializer)
+        return Response(serializer.data, status=HTTP_201_CREATED)
 
-        # Authenticate the user account based on the provided credentials
-        username = serializer.data.get('username')
-        password = serializer.data.get('password')
-        if not username or not password:
-            raise AuthenticationFailed("Username and password must be provided to provision a token.")
-        user = authenticate(request=request, username=username, password=password)
-        if user is None:
-            raise AuthenticationFailed("Invalid username/password")
-
-        # Create a new Token for the User
-        token = Token(user=user)
-        token.save()
-        data = serializers.TokenSerializer(token, context={'request': request}).data
-        # Manually append the token key, which is normally write-only
-        data['key'] = token.key
-
-        return Response(data, status=HTTP_201_CREATED)
-
-    def get_serializer_class(self):
-        return serializers.TokenSerializer
+    def perform_create(self, serializer):
+        model = serializer.Meta.model
+        logger = logging.getLogger(f'netbox.api.views.TokenProvisionView')
+        logger.info(f"Creating new {model._meta.verbose_name}")
+        serializer.save()
 
 
 #