perf(authentication): 优化connection token的使用

apps/authentication/api/auth.py

@@ -2,54 +2,59 @@
 #
 import uuid
 
+from django.conf import settings
 from django.core.cache import cache
-from django.shortcuts import get_object_or_404
 from rest_framework.response import Response
-from rest_framework.views import APIView
+from rest_framework.generics import CreateAPIView
 
 from common.utils import get_logger
-from common.permissions import IsOrgAdminOrAppUser
+from common.permissions import IsSuperUserOrAppUser
 from orgs.mixins.api import RootOrgViewMixin
-from users.models import User
-from assets.models import Asset, SystemUser
+
+from ..serializers import ConnectionTokenSerializer
 
 logger = get_logger(__name__)
-__all__ = [
-    'UserConnectionTokenApi',
-]
+__all__ = ['UserConnectionTokenApi']
 
 
-class UserConnectionTokenApi(RootOrgViewMixin, APIView):
-    permission_classes = (IsOrgAdminOrAppUser,)
+class UserConnectionTokenApi(RootOrgViewMixin, CreateAPIView):
+    permission_classes = (IsSuperUserOrAppUser,)
+    serializer_class = ConnectionTokenSerializer
 
-    def post(self, request):
-        user_id = request.data.get('user', '')
-        asset_id = request.data.get('asset', '')
-        system_user_id = request.data.get('system_user', '')
+    def perform_create(self, serializer):
+        user = serializer.validated_data['user']
+        asset = serializer.validated_data['asset']
+        system_user = serializer.validated_data['system_user']
         token = str(uuid.uuid4())
-        user = get_object_or_404(User, id=user_id)
-        asset = get_object_or_404(Asset, id=asset_id)
-        system_user = get_object_or_404(SystemUser, id=system_user_id)
         value = {
-            'user': user_id,
+            'user': str(user.id),
             'username': user.username,
-            'asset': asset_id,
+            'asset': str(asset.id),
             'hostname': asset.hostname,
-            'system_user': system_user_id,
+            'system_user': str(system_user.id),
             'system_user_name': system_user.name
         }
         cache.set(token, value, timeout=20)
+        return token
+
+    def create(self, request, *args, **kwargs):
+        if not settings.CONNECTION_TOKEN_ENABLED:
+            data = {'error': 'Connection token disabled'}
+            return Response(data, status=400)
+
+        if not request.user.is_superuser:
+            data = {'error': 'Only super user can create token'}
+            return Response(data, status=403)
+
+        serializer = self.get_serializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+        token = self.perform_create(serializer)
         return Response({"token": token}, status=201)
 
     def get(self, request):
         token = request.query_params.get('token')
-        user_only = request.query_params.get('user-only', None)
         value = cache.get(token, None)
 
         if not value:
             return Response('', status=404)
-
-        if not user_only:
-            return Response(value)
-        else:
-            return Response({'user': value['user']})
+        return Response(value)

apps/authentication/serializers.py

@@ -11,6 +11,7 @@ from .models import AccessKey, LoginConfirmSetting, SSOToken
 __all__ = [
     'AccessKeySerializer', 'OtpVerifySerializer', 'BearerTokenSerializer',
     'MFAChallengeSerializer', 'LoginConfirmSettingSerializer', 'SSOTokenSerializer',
+    'ConnectionTokenSerializer',
 ]
 
 
@@ -82,3 +83,33 @@ class SSOTokenSerializer(serializers.Serializer):
     username = serializers.CharField(write_only=True)
     login_url = serializers.CharField(read_only=True)
     next = serializers.CharField(write_only=True, allow_blank=True, required=False, allow_null=True)
+
+
+class ConnectionTokenSerializer(serializers.Serializer):
+    user = serializers.CharField(max_length=128, required=True)
+    system_user = serializers.CharField(max_length=128, required=True)
+    asset = serializers.CharField(max_length=128, required=True)
+
+    @staticmethod
+    def validate_user(user_id):
+        from users.models import User
+        user = User.objects.filter(id=user_id).first()
+        if user is None:
+            raise serializers.ValidationError('user id not exist')
+        return user
+
+    @staticmethod
+    def validate_system_user(system_user_id):
+        from assets.models import SystemUser
+        system_user = SystemUser.objects.filter(id=system_user_id).first()
+        if system_user is None:
+            raise serializers.ValidationError('system_user id not exist')
+        return system_user
+
+    @staticmethod
+    def validate_asset(asset_id):
+        from assets.models import Asset
+        asset = Asset.objects.filter(id=asset_id).first()
+        if asset is None:
+            raise serializers.ValidationError('asset id not exist')
+        return asset

apps/jumpserver/conf.py

@@ -280,7 +280,8 @@ class Config(dict):
         'SESSION_COOKIE_SECURE': False,
         'CSRF_COOKIE_SECURE': False,
         'REFERER_CHECK_ENABLED': False,
-        'SERVER_REPLAY_STORAGE': {}
+        'SERVER_REPLAY_STORAGE': {},
+        'CONNECTION_TOKEN_ENABLED': False,
     }
 
     def compatible_auth_openid_of_key(self):

apps/jumpserver/settings/custom.py

@@ -115,3 +115,5 @@ DATETIME_DISPLAY_FORMAT = '%Y-%m-%d %H:%M:%S'
 
 TICKETS_ENABLED = CONFIG.TICKETS_ENABLED
 REFERER_CHECK_ENABLED = CONFIG.REFERER_CHECK_ENABLED
+
+CONNECTION_TOKEN_ENABLED = CONFIG.CONNECTION_TOKEN_ENABLED