perf: 修改支持客户端拉起

perf: remove print

apps/authentication/api/connection_token.py

@@ -1,7 +1,9 @@
 # -*- coding: utf-8 -*-
 #
 import urllib.parse
+import json
 import base64
+from typing import Callable
 
 from django.conf import settings
 from django.core.cache import cache
@@ -9,6 +11,7 @@ from django.shortcuts import get_object_or_404
 from django.http import HttpResponse
 from django.utils.translation import ugettext as _
 from rest_framework.response import Response
+from rest_framework.request import Request
 from rest_framework.viewsets import GenericViewSet
 from rest_framework.decorators import action
 from rest_framework.exceptions import PermissionDenied
@@ -18,71 +21,32 @@ from authentication.signals import post_auth_failed, post_auth_success
 from common.utils import get_logger, random_string
 from common.drf.api import SerializerMixin
 from common.permissions import IsSuperUserOrAppUser, IsValidUser, IsSuperUser
-
 from orgs.mixins.api import RootOrgViewMixin
 
 from ..serializers import (
     ConnectionTokenSerializer, ConnectionTokenSecretSerializer,
-    RDPFileSerializer
 )
 
 logger = get_logger(__name__)
 __all__ = ['UserConnectionTokenViewSet']
 
 
-class UserConnectionTokenViewSet(RootOrgViewMixin, SerializerMixin, GenericViewSet):
+class ClientProtocolMixin:
-    permission_classes = (IsSuperUserOrAppUser,)
+    request: Request
-    serializer_classes = {
+    get_serializer: Callable
-        'default': ConnectionTokenSerializer,
+    create_token: Callable
-        'get_secret_detail': ConnectionTokenSecretSerializer,
-        'get_rdp_file': RDPFileSerializer
-    }
-    CACHE_KEY_PREFIX = 'CONNECTION_TOKEN_{}'
 
-    @staticmethod
+    def get_request_resource(self, serializer):
-    def check_resource_permission(user, asset, application, system_user):
+        asset = serializer.validated_data.get('asset')
-        from perms.utils.asset import has_asset_system_permission
+        application = serializer.validated_data.get('application')
-        from perms.utils.application import has_application_system_permission
+        system_user = serializer.validated_data['system_user']
-        if asset and not has_asset_system_permission(user, asset, system_user):
-            error = f'User not has this asset and system user permission: ' \
-                    f'user={user.id} system_user={system_user.id} asset={asset.id}'
-            raise PermissionDenied(error)
-        if application and not has_application_system_permission(user, application, system_user):
-            error = f'User not has this application and system user permission: ' \
-                    f'user={user.id} system_user={system_user.id} application={application.id}'
-            raise PermissionDenied(error)
-        return True
 
-    def create_token(self, user, asset, application, system_user, ttl=5 * 60):
+        user = serializer.validated_data.get('user')
-        if not self.request.user.is_superuser and user != self.request.user:
+        if not user or not self.request.user.is_superuser:
-            raise PermissionDenied('Only super user can create user token')
+            user = self.request.user
-        self.check_resource_permission(user, asset, application, system_user)
+        return asset, application, system_user, user
-        token = random_string(36)
-        value = {
-            'user': str(user.id),
-            'username': user.username,
-            'system_user': str(system_user.id),
-            'system_user_name': system_user.name
-        }
 
-        if asset:
+    def get_rdp_file_content(self, serializer):
-            value.update({
-                'type': 'asset',
-                'asset': str(asset.id),
-                'hostname': asset.hostname,
-            })
-        elif application:
-            value.update({
-                'type': 'application',
-                'application': application.id,
-                'application_name': str(application)
-            })
-
-        key = self.CACHE_KEY_PREFIX.format(token)
-        cache.set(key, value, timeout=ttl)
-        return token
-
-    def create_rdp_file(self):
         options = {
             'full address:s': '',
             'username:s': '',
@@ -114,24 +78,15 @@ class UserConnectionTokenViewSet(RootOrgViewMixin, SerializerMixin, GenericViewS
             # 'remoteapplicationname:s': 'Firefox',
             # 'remoteapplicationcmdline:s': '',
         }
-        if self.request.method == 'GET':
-            data = self.request.query_params
-        else:
-            data = request.data
-        serializer = self.get_serializer(data=data)
-        serializer.is_valid(raise_exception=True)
 
-        asset = serializer.validated_data.get('asset')
+        asset, application, system_user, user = self.get_request_resource(serializer)
-        application = serializer.validated_data.get('application')
+        height = self.request.query_params.get('height')
-        system_user = serializer.validated_data['system_user']
+        width = self.request.query_params.get('width')
-        height = serializer.validated_data.get('height')
-        width = serializer.validated_data.get('width')
-        user = self.request.user
         token = self.create_token(user, asset, application, system_user)
 
         address = settings.TERMINAL_RDP_ADDR
         if not address or address == 'localhost:3389':
-            address = request.get_host().split(':')[0] + ':3389'
+            address = self.request.get_host().split(':')[0] + ':3389'
         options['full address:s'] = address
         options['username:s'] = '{}|{}'.format(user.username, token)
         if system_user.ad_domain:
@@ -152,30 +107,61 @@ class UserConnectionTokenViewSet(RootOrgViewMixin, SerializerMixin, GenericViewS
             name = '*'
         return name, content
 
-    def create(self, request, *args, **kwargs):
-        serializer = self.get_serializer(data=request.data)
-        serializer.is_valid(raise_exception=True)
-
-        asset = serializer.validated_data.get('asset')
-        application = serializer.validated_data.get('application')
-        system_user = serializer.validated_data['system_user']
-        user = serializer.validated_data.get('user')
-        token = self.create_token(user, asset, application, system_user)
-        return Response({"token": token}, status=201)
-
     @action(methods=['POST', 'GET'], detail=False, url_path='rdp/file', permission_classes=[IsValidUser])
     def get_rdp_file(self, request, *args, **kwargs):
-        name, data = self.create_rdp_file()
+        if self.request.method == 'GET':
+            data = self.request.query_params
+        else:
+            data = self.request.data
+        serializer = self.get_serializer(data=data)
+        serializer.is_valid(raise_exception=True)
+        name, data = self.get_rdp_file_content(serializer)
         response = HttpResponse(data, content_type='application/octet-stream')
         filename = "{}-{}-jumpserver.rdp".format(self.request.user.username, name)
         filename = urllib.parse.quote(filename)
         response['Content-Disposition'] = 'attachment; filename*=UTF-8\'\'%s' % filename
         return response
 
-    @action(methods=['POST', 'GET'], detail=False, url_path='rdp/rouse', permission_classes=[IsValidUser])
+    def get_valid_serializer(self):
-    def get_rdp_rouse(self, request, *args, **kwargs):
+        if self.request.method == 'GET':
-        data = self.create_rdp_file()[1]
+            data = self.request.query_params
-        return Response(data=dict(data=base64.b64encode(data.encode())))
+        else:
+            data = self.request.data
+        serializer = self.get_serializer(data=data)
+        serializer.is_valid(raise_exception=True)
+        return serializer
+
+    def get_client_protocol_data(self, serializer):
+        asset, application, system_user, user = self.get_request_resource(serializer)
+        protocol = system_user.protocol
+        if protocol == 'rdp':
+            name, config = self.get_rdp_file_content(serializer)
+        elif protocol == 'vnc':
+            raise HttpResponse(status=404, data={"error": "VNC not support"})
+        else:
+            config = 'ssh://system_user@asset@user@jumpserver-ssh'
+        data = {
+            "protocol": system_user.protocol,
+            "username": user.username,
+            "config": config
+        }
+        return data
+
+    @action(methods=['POST', 'GET'], detail=False, url_path='client-url', permission_classes=[IsValidUser])
+    def get_client_protocol_url(self, request, *args, **kwargs):
+        serializer = self.get_valid_serializer()
+        protocol_data = self.get_client_protocol_data(serializer)
+        protocol_data = base64.b64encode(json.dumps(protocol_data).encode()).decode()
+        data = {
+            'url': 'jms://{}'.format(protocol_data),
+        }
+        return Response(data=data)
+
+
+class SecretDetailMixin:
+    valid_token: Callable
+    request: Request
+    get_serializer: Callable
 
     @staticmethod
     def _get_application_secret_detail(application):
@@ -222,6 +208,100 @@ class UserConnectionTokenViewSet(RootOrgViewMixin, SerializerMixin, GenericViewS
             'actions': actions,
         }
 
+    @action(methods=['POST'], detail=False, permission_classes=[IsSuperUserOrAppUser], url_path='secret-info/detail')
+    def get_secret_detail(self, request, *args, **kwargs):
+        token = request.data.get('token', '')
+        try:
+            value, user, system_user, asset, app, expired_at = self.valid_token(token)
+        except serializers.ValidationError as e:
+            post_auth_failed.send(
+                sender=self.__class__, username='', request=self.request,
+                reason=_('Invalid token')
+            )
+            raise e
+
+        data = dict(user=user, system_user=system_user, expired_at=expired_at)
+        if asset:
+            asset_detail = self._get_asset_secret_detail(asset, user=user, system_user=system_user)
+            system_user.load_asset_more_auth(asset.id, user.username, user.id)
+            data['type'] = 'asset'
+            data.update(asset_detail)
+        else:
+            app_detail = self._get_application_secret_detail(app)
+            system_user.load_app_more_auth(app.id, user.id)
+            data['type'] = 'application'
+            data.update(app_detail)
+
+        self.request.session['auth_backend'] = settings.AUTH_BACKEND_AUTH_TOKEN
+        post_auth_success.send(sender=self.__class__, user=user, request=self.request, login_type='T')
+
+        serializer = self.get_serializer(data)
+        return Response(data=serializer.data, status=200)
+
+
+class UserConnectionTokenViewSet(
+    RootOrgViewMixin, SerializerMixin, ClientProtocolMixin,
+    SecretDetailMixin, GenericViewSet
+):
+    permission_classes = (IsSuperUserOrAppUser,)
+    serializer_classes = {
+        'default': ConnectionTokenSerializer,
+        'get_secret_detail': ConnectionTokenSecretSerializer,
+    }
+    CACHE_KEY_PREFIX = 'CONNECTION_TOKEN_{}'
+
+    @staticmethod
+    def check_resource_permission(user, asset, application, system_user):
+        from perms.utils.asset import has_asset_system_permission
+        from perms.utils.application import has_application_system_permission
+
+        if asset and not has_asset_system_permission(user, asset, system_user):
+            error = f'User not has this asset and system user permission: ' \
+                    f'user={user.id} system_user={system_user.id} asset={asset.id}'
+            raise PermissionDenied(error)
+        if application and not has_application_system_permission(user, application, system_user):
+            error = f'User not has this application and system user permission: ' \
+                    f'user={user.id} system_user={system_user.id} application={application.id}'
+            raise PermissionDenied(error)
+        return True
+
+    def create_token(self, user, asset, application, system_user, ttl=5 * 60):
+        if not self.request.user.is_superuser and user != self.request.user:
+            raise PermissionDenied('Only super user can create user token')
+        self.check_resource_permission(user, asset, application, system_user)
+        token = random_string(36)
+        value = {
+            'user': str(user.id),
+            'username': user.username,
+            'system_user': str(system_user.id),
+            'system_user_name': system_user.name
+        }
+
+        if asset:
+            value.update({
+                'type': 'asset',
+                'asset': str(asset.id),
+                'hostname': asset.hostname,
+            })
+        elif application:
+            value.update({
+                'type': 'application',
+                'application': application.id,
+                'application_name': str(application)
+            })
+
+        key = self.CACHE_KEY_PREFIX.format(token)
+        cache.set(key, value, timeout=ttl)
+        return token
+
+    def create(self, request, *args, **kwargs):
+        serializer = self.get_serializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+
+        asset, application, system_user, user = self.get_request_resource(serializer)
+        token = self.create_token(user, asset, application, system_user)
+        return Response({"token": token}, status=201)
+
     def valid_token(self, token):
         from users.models import User
         from assets.models import SystemUser, Asset
@@ -254,39 +334,8 @@ class UserConnectionTokenViewSet(RootOrgViewMixin, SerializerMixin, GenericViewS
 
         if not has_perm:
             raise serializers.ValidationError('Permission expired or invalid')
-
         return value, user, system_user, asset, app, expired_at
 
-    @action(methods=['POST'], detail=False, permission_classes=[IsSuperUserOrAppUser], url_path='secret-info/detail')
-    def get_secret_detail(self, request, *args, **kwargs):
-        token = request.data.get('token', '')
-        try:
-            value, user, system_user, asset, app, expired_at = self.valid_token(token)
-        except serializers.ValidationError as e:
-            post_auth_failed.send(
-                sender=self.__class__, username='', request=self.request,
-                reason=_('Invalid token')
-            )
-            raise e
-
-        data = dict(user=user, system_user=system_user, expired_at=expired_at)
-        if asset:
-            asset_detail = self._get_asset_secret_detail(asset, user=user, system_user=system_user)
-            system_user.load_asset_more_auth(asset.id, user.username, user.id)
-            data['type'] = 'asset'
-            data.update(asset_detail)
-        else:
-            app_detail = self._get_application_secret_detail(app)
-            system_user.load_app_more_auth(app.id, user.id)
-            data['type'] = 'application'
-            data.update(app_detail)
-
-        self.request.session['auth_backend'] = settings.AUTH_BACKEND_AUTH_TOKEN
-        post_auth_success.send(sender=self.__class__, user=user, request=self.request, login_type='T')
-
-        serializer = self.get_serializer(data)
-        return Response(data=serializer.data, status=200)
-
     def get_permissions(self):
         if self.action in ["create", "get_rdp_file"]:
             if self.request.data.get('user', None):

apps/authentication/serializers.py

@@ -198,7 +198,3 @@ class ConnectionTokenSecretSerializer(serializers.Serializer):
     actions = ActionsField()
     expired_at = serializers.IntegerField()
 
-
-class RDPFileSerializer(ConnectionTokenSerializer):
-    width = serializers.IntegerField(allow_null=True, max_value=3112, min_value=100, required=False)
-    height = serializers.IntegerField(allow_null=True, max_value=4096, min_value=100, required=False)