perf: merge connect token rdp option

apps/authentication/api/connection_token.py

@@ -19,12 +19,12 @@ from common.utils import random_string
 from common.utils.django import get_request_os
 from orgs.mixins.api import RootOrgViewMixin
 from perms.models import ActionChoices
-from terminal.const import NativeClient, TerminalType
+from terminal.connect_methods import NativeClient, ConnectMethodUtil
 from terminal.models import EndpointRule, Applet
 from ..models import ConnectionToken
 from ..serializers import (
     ConnectionTokenSerializer, ConnectionTokenSecretSerializer,
-    SuperConnectionTokenSerializer,
+    SuperConnectionTokenSerializer, ConnectTokenAppletOptionSerializer
 )
 
 __all__ = ['ConnectionTokenViewSet', 'SuperConnectionTokenViewSet']
@@ -115,7 +115,8 @@ class RDPFileClientProtocolURLMixin:
         rdp_options['audiomode:i'] = self.parse_env_bool('JUMPSERVER_DISABLE_AUDIO', 'false', '2', '0')
 
         # 设置远程应用
-        self.set_applet_info(token, rdp_options)
+        remote_app_options = token.get_remote_app_option()
+        rdp_options.update(remote_app_options)
 
         # 文件名
         name = token.asset.name
@@ -145,7 +146,7 @@ class RDPFileClientProtocolURLMixin:
         _os = get_request_os(self.request)
 
         connect_method_name = token.connect_method
-        connect_method_dict = TerminalType.get_connect_method(
+        connect_method_dict = ConnectMethodUtil.get_connect_method(
             token.connect_method, token.protocol, _os
         )
         if connect_method_dict is None:
@@ -227,38 +228,16 @@ class ConnectionTokenViewSet(ExtraActionApiMixin, RootOrgViewMixin, JMSModelView
     search_fields = filterset_fields
     serializer_classes = {
         'default': ConnectionTokenSerializer,
-        'get_secret_detail': ConnectionTokenSecretSerializer,
     }
     rbac_perms = {
         'list': 'authentication.view_connectiontoken',
         'retrieve': 'authentication.view_connectiontoken',
         'create': 'authentication.add_connectiontoken',
         'expire': 'authentication.add_connectiontoken',
-        'get_secret_detail': 'authentication.view_connectiontokensecret',
         'get_rdp_file': 'authentication.add_connectiontoken',
         'get_client_protocol_url': 'authentication.add_connectiontoken',
     }
 
-    @action(methods=['POST'], detail=False, url_path='secret')
-    def get_secret_detail(self, request, *args, **kwargs):
-        """ 非常重要的 api, 在逻辑层再判断一下 rbac 权限, 双重保险 """
-        rbac_perm = 'authentication.view_connectiontokensecret'
-        if not request.user.has_perm(rbac_perm):
-            raise PermissionDenied('Not allow to view secret')
-
-        token_id = request.data.get('id') or ''
-        token = get_object_or_404(ConnectionToken, pk=token_id)
-        if token.is_expired:
-            raise ValidationError({'id': 'Token is expired'})
-
-        token.is_valid()
-        serializer = self.get_serializer(instance=token)
-        expire_now = request.data.get('expire_now', True)
-        if expire_now:
-            token.expire()
-
-        return Response(serializer.data, status=status.HTTP_200_OK)
-
     def get_queryset(self):
         queryset = ConnectionToken.objects \
             .filter(user=self.request.user) \
@@ -305,10 +284,14 @@ class ConnectionTokenViewSet(ExtraActionApiMixin, RootOrgViewMixin, JMSModelView
 class SuperConnectionTokenViewSet(ConnectionTokenViewSet):
     serializer_classes = {
         'default': SuperConnectionTokenSerializer,
+        'get_secret_detail': ConnectionTokenSecretSerializer,
     }
     rbac_perms = {
         'create': 'authentication.add_superconnectiontoken',
-        'renewal': 'authentication.add_superconnectiontoken'
+        'renewal': 'authentication.add_superconnectiontoken',
+        'get_secret_detail': 'authentication.view_connectiontokensecret',
+        'get_applet_info': 'authentication.view_superconnectiontoken',
+        'release_applet_account': 'authentication.view_superconnectiontoken',
     }
 
     def get_queryset(self):
@@ -332,3 +315,38 @@ class SuperConnectionTokenViewSet(ConnectionTokenViewSet):
             'msg': f'Token is renewed, date expired: {date_expired}'
         }
         return Response(data=data, status=status.HTTP_200_OK)
+
+    @action(methods=['POST'], detail=False, url_path='secret')
+    def get_secret_detail(self, request, *args, **kwargs):
+        """ 非常重要的 api, 在逻辑层再判断一下 rbac 权限, 双重保险 """
+        rbac_perm = 'authentication.view_connectiontokensecret'
+        if not request.user.has_perm(rbac_perm):
+            raise PermissionDenied('Not allow to view secret')
+
+        token_id = request.data.get('id') or ''
+        token = get_object_or_404(ConnectionToken, pk=token_id)
+        if token.is_expired:
+            raise ValidationError({'id': 'Token is expired'})
+
+        token.is_valid()
+        serializer = self.get_serializer(instance=token)
+        expire_now = request.data.get('expire_now', True)
+        if expire_now:
+            token.expire()
+        return Response(serializer.data, status=status.HTTP_200_OK)
+
+    @action(methods=['POST'], detail=False, url_path='applet-option')
+    def get_applet_info(self, *args, **kwargs):
+        token_id = self.request.data.get('id')
+        token = get_object_or_404(ConnectionToken, pk=token_id)
+        if token.is_expired:
+            return Response({'error': 'Token expired'}, status=status.HTTP_400_BAD_REQUEST)
+        data = token.get_applet_option()
+        serializer = ConnectTokenAppletOptionSerializer(data)
+        return Response(serializer.data)
+
+    @action(methods=['DELETE', 'POST'], detail=False, url_path='applet-account/release')
+    def release_applet_account(self, *args, **kwargs):
+        account_id = self.request.data.get('id')
+        msg = ConnectionToken.release_applet_account(account_id)
+        return Response({'msg': msg})

apps/authentication/models/connection_token.py

@@ -1,6 +1,9 @@
+import base64
+import json
 from datetime import timedelta
 
 from django.conf import settings
+from django.core.cache import cache
 from django.db import models
 from django.utils import timezone
 from django.utils.translation import ugettext_lazy as _
@@ -9,9 +12,10 @@ from rest_framework.exceptions import PermissionDenied
 from assets.const import Protocol
 from common.db.fields import EncryptCharField
 from common.db.models import JMSBaseModel
-from common.utils import lazyproperty, pretty_string
+from common.utils import lazyproperty, pretty_string, bulk_get
 from common.utils.timezone import as_current_tz
 from orgs.mixins.models import OrgModelMixin
+from terminal.models import Applet
 
 
 def date_expired_default():
@@ -101,6 +105,9 @@ class ConnectionToken(OrgModelMixin, JMSBaseModel):
             error = _('No account')
             raise PermissionDenied(error)
 
+        if timezone.now() - self.date_created < timedelta(seconds=60):
+            return True, None
+
         if not self.permed_account or not self.permed_account.actions:
             msg = 'user `{}` not has asset `{}` permission for login `{}`'.format(
                 self.user, self.asset, self.account
@@ -115,6 +122,75 @@ class ConnectionToken(OrgModelMixin, JMSBaseModel):
     def platform(self):
         return self.asset.platform
 
+    @lazyproperty
+    def connect_method_object(self):
+        from common.utils import get_request_os
+        from jumpserver.utils import get_current_request
+        from terminal.connect_methods import ConnectMethodUtil
+
+        request = get_current_request()
+        os = get_request_os(request) if request else 'windows'
+        method = ConnectMethodUtil.get_connect_method(
+            self.connect_method, protocol=self.protocol, os=os
+        )
+        return method
+
+    def get_remote_app_option(self):
+        cmdline = {
+            'app_name': self.connect_method,
+            'user_id': str(self.user.id),
+            'asset_id': str(self.asset.id),
+            'token_id': str(self.id)
+        }
+        cmdline_b64 = base64.b64encode(json.dumps(cmdline).encode()).decode()
+        app = '||tinker'
+        options = {
+            'remoteapplicationmode:i': '1',
+            'remoteapplicationprogram:s': app,
+            'remoteapplicationname:s': app,
+            'alternate shell:s': app,
+            'remoteapplicationcmdline:s': cmdline_b64,
+        }
+        return options
+
+    def get_applet_option(self):
+        method = self.connect_method_object
+        if not method or method.get('type') != 'applet' or method.get('disabled', False):
+            return None
+
+        applet = Applet.objects.filter(name=method.get('value')).first()
+        if not applet:
+            return None
+
+        host_account = applet.select_host_account()
+        if not host_account:
+            return None
+
+        host, account, lock_key, ttl = bulk_get(host_account, ('host', 'account', 'lock_key', 'ttl'))
+        gateway = host.gateway.select_gateway() if host.domain else None
+
+        data = {
+            'id': account.id,
+            'applet': applet,
+            'host': host,
+            'gateway': gateway,
+            'account': account,
+            'remote_app_option': self.get_remote_app_option()
+        }
+        token_account_relate_key = f'token_account_relate_{account.id}'
+        cache.set(token_account_relate_key, lock_key, ttl)
+        return data
+
+    @staticmethod
+    def release_applet_account(account_id):
+        token_account_relate_key = f'token_account_relate_{account_id}'
+        lock_key = cache.get(token_account_relate_key)
+        if lock_key:
+            cache.delete(lock_key)
+            cache.delete(token_account_relate_key)
+            return 'released'
+        return 'not found or expired'
+
     @lazyproperty
     def account_object(self):
         from assets.models import Account

apps/authentication/serializers/connect_token_secret.py

@@ -1,19 +1,17 @@
 from django.utils.translation import ugettext_lazy as _
 from rest_framework import serializers
 
-from common.drf.fields import ObjectRelatedField
 from acls.models import CommandGroup, CommandFilterACL
 from assets.models import Asset, Account, Platform, Gateway, Domain
 from assets.serializers import PlatformSerializer, AssetProtocolsSerializer
-from users.models import User
-from perms.serializers.permission import ActionChoicesField
+from common.drf.fields import ObjectRelatedField
 from orgs.mixins.serializers import OrgResourceModelSerializerMixin
-
+from perms.serializers.permission import ActionChoicesField
+from users.models import User
 from ..models import ConnectionToken
 
-
 __all__ = [
-    'ConnectionTokenSecretSerializer',
+    'ConnectionTokenSecretSerializer', 'ConnectTokenAppletOptionSerializer'
 ]
 
 
@@ -96,6 +94,24 @@ class _ConnectionTokenPlatformSerializer(PlatformSerializer):
         return names
 
 
+class _ConnectionTokenConnectMethodSerializer(serializers.Serializer):
+    name = serializers.CharField(label=_('Name'))
+    protocol = serializers.CharField(label=_('Protocol'))
+    os = serializers.CharField(label=_('OS'))
+    is_builtin = serializers.BooleanField(label=_('Is builtin'))
+    is_active = serializers.BooleanField(label=_('Is active'))
+    platform = _ConnectionTokenPlatformSerializer(label=_('Platform'))
+    action = ActionChoicesField(label=_('Action'))
+    options = serializers.JSONField(label=_('Options'))
+
+
+class _ConnectTokenConnectMethodSerializer(serializers.Serializer):
+    label = serializers.CharField(label=_('Label'))
+    value = serializers.CharField(label=_('Value'))
+    type = serializers.CharField(label=_('Type'))
+    component = serializers.CharField(label=_('Component'))
+
+
 class ConnectionTokenSecretSerializer(OrgResourceModelSerializerMixin):
     user = _ConnectionTokenUserSerializer(read_only=True)
     asset = _ConnectionTokenAssetSerializer(read_only=True)
@@ -104,30 +120,28 @@ class ConnectionTokenSecretSerializer(OrgResourceModelSerializerMixin):
     platform = _ConnectionTokenPlatformSerializer(read_only=True)
     domain = ObjectRelatedField(queryset=Domain.objects, required=False, label=_('Domain'))
     command_filter_acls = _ConnectionTokenCommandFilterACLSerializer(read_only=True, many=True)
+    expire_now = serializers.BooleanField(label=_('Expired now'), write_only=True, default=True)
+    connect_method = _ConnectTokenConnectMethodSerializer(read_only=True, source='connect_method_object')
     actions = ActionChoicesField()
     expire_at = serializers.IntegerField()
-    expire_now = serializers.BooleanField(label=_('Expired now'), write_only=True, default=True)
-    connect_method = serializers.SerializerMethodField(label=_('Connect method'))
 
     class Meta:
         model = ConnectionToken
         fields = [
             'id', 'value', 'user', 'asset', 'account',
             'platform', 'command_filter_acls', 'protocol',
-            'domain', 'gateway', 'actions', 'expire_at', 'expire_now',
-            'connect_method'
+            'domain', 'gateway', 'actions', 'expire_at',
+            'expire_now', 'connect_method',
         ]
         extra_kwargs = {
             'value': {'read_only': True},
         }
 
-    def get_connect_method(self, obj):
-        from terminal.const import TerminalType
-        from common.utils import get_request_os
-        request = self.context.get('request')
-        if request:
-            os = get_request_os(request)
-        else:
-            os = 'windows'
-        method = TerminalType.get_connect_method(obj.connect_method, protocol=obj.protocol, os=os)
-        return method
+
+class ConnectTokenAppletOptionSerializer(serializers.Serializer):
+    id = serializers.CharField(label=_('ID'))
+    applet = ObjectRelatedField(read_only=True)
+    host = _ConnectionTokenAssetSerializer(read_only=True)
+    account = _ConnectionTokenAccountSerializer(read_only=True)
+    gateway = _ConnectionTokenGatewaySerializer(read_only=True)
+    remote_app_option = serializers.JSONField(read_only=True)

apps/authentication/serializers/connection_token.py

@@ -2,7 +2,6 @@ from django.utils.translation import ugettext_lazy as _
 from rest_framework import serializers
 
 from orgs.mixins.serializers import OrgResourceModelSerializerMixin
-
 from ..models import ConnectionToken
 
 __all__ = [