Просмотр исходного кода

return v2 token for API Token create

Arthur 2 недель назад
Родитель
Сommit
ae811a07fb
2 измененных файлов с 71 добавлено и 0 удалено
  1. 24 0
      netbox/users/api/views.py
  2. 47 0
      netbox/users/tests/test_api.py

+ 24 - 0
netbox/users/api/views.py

@@ -52,6 +52,30 @@ class TokenViewSet(NetBoxModelViewSet):
     serializer_class = serializers.TokenSerializer
     filterset_class = filtersets.TokenFilterSet
 
+    def create(self, request, *args, **kwargs):
+        # The plaintext token value is held only in memory after creation; for v2 tokens the database
+        # stores only an HMAC digest, so it cannot be recovered later. Preserve it across the re-fetch
+        # performed in the parent class so the response returned to the client includes a usable token.
+        serializer = self.get_serializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+        bulk_create = getattr(serializer, 'many', False)
+        self.perform_create(serializer)
+
+        if bulk_create:
+            plaintext_by_pk = {obj.pk: obj.token for obj in serializer.instance}
+            qs = list(self.get_queryset().filter(pk__in=plaintext_by_pk.keys()).order_by('pk'))
+            for obj in qs:
+                obj._token = plaintext_by_pk.get(obj.pk)
+        else:
+            plaintext = serializer.instance.token
+            qs = self.get_queryset().get(pk=serializer.instance.pk)
+            qs._token = plaintext
+
+        serializer = self.get_serializer(qs, many=bulk_create)
+
+        headers = self.get_success_headers(serializer.data)
+        return Response(serializer.data, status=HTTP_201_CREATED, headers=headers)
+
 
 class TokenProvisionView(APIView):
     """

+ 47 - 0
netbox/users/tests/test_api.py

@@ -329,6 +329,53 @@ class TokenTest(
         token1.refresh_from_db()
         self.assertEqual(token1.user, user1, "Token's user should not have changed")
 
+    def test_create_token_returns_plaintext(self):
+        """
+        Test that creating a Token via the REST API returns the usable plaintext value
+        in the response. For v2 tokens this value cannot be recovered later because the
+        database stores only an HMAC digest.
+        """
+        self.add_permissions('users.add_token')
+        user = User.objects.create_user(username='token_plaintext_user')
+        url = reverse('users-api:token-list')
+
+        response = self.client.post(url, {'user': user.pk}, format='json', **self.header)
+        self.assertEqual(response.status_code, 201)
+        self.assertIsNotNone(response.data['token'])
+        self.assertEqual(len(response.data['token']), TOKEN_DEFAULT_LENGTH)
+
+        # The returned plaintext must authenticate against the stored token
+        token = Token.objects.get(pk=response.data['id'])
+        self.assertTrue(token.validate(response.data['token']))
+
+    def test_bulk_create_tokens_returns_plaintexts(self):
+        """
+        Test that bulk-creating Tokens via the REST API returns the plaintext value for
+        each created Token in the response.
+        """
+        self.add_permissions('users.add_token')
+        users = [
+            User.objects.create_user(username='token_bulk_user1'),
+            User.objects.create_user(username='token_bulk_user2'),
+        ]
+        data = [{'user': u.pk} for u in users]
+        url = reverse('users-api:token-list')
+
+        response = self.client.post(url, data, format='json', **self.header)
+        self.assertEqual(response.status_code, 201)
+        self.assertEqual(len(response.data), len(data))
+
+        plaintexts = set()
+        for obj in response.data:
+            self.assertIsNotNone(obj['token'])
+            self.assertEqual(len(obj['token']), TOKEN_DEFAULT_LENGTH)
+            plaintexts.add(obj['token'])
+            token = Token.objects.get(pk=obj['id'])
+            self.assertTrue(token.validate(obj['token']))
+
+        # Each token should be unique
+        self.assertEqual(len(plaintexts), len(data))
+
 
 class ObjectPermissionTest(
     # No GraphQL support for ObjectPermission