From 24da1e7d9178a6de5e4f5baa8fea2139c32983cf Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Tue, 20 Dec 2022 11:05:00 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BF=AE=E6=94=B9=20connect=20token,?=
 =?UTF-8?q?=20=E8=8E=B7=E5=8F=96=20applet=20info?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/authentication/api/connection_token.py | 37 +++++----------------
 apps/terminal/connect_methods.py            |  6 ++--
 2 files changed, 11 insertions(+), 32 deletions(-)

diff --git a/apps/authentication/api/connection_token.py b/apps/authentication/api/connection_token.py
index c77a04fa7..782346fc5 100644
--- a/apps/authentication/api/connection_token.py
+++ b/apps/authentication/api/connection_token.py
@@ -20,7 +20,7 @@ from common.utils.django import get_request_os
 from orgs.mixins.api import RootOrgViewMixin
 from perms.models import ActionChoices
 from terminal.connect_methods import NativeClient, ConnectMethodUtil
-from terminal.models import EndpointRule, Applet
+from terminal.models import EndpointRule
 from ..models import ConnectionToken
 from ..serializers import (
     ConnectionTokenSerializer, ConnectionTokenSecretSerializer,
@@ -34,30 +34,6 @@ class RDPFileClientProtocolURLMixin:
     request: Request
     get_serializer: callable
 
-    @staticmethod
-    def set_applet_info(token, rdp_options):
-        # remote-app
-        applet = Applet.objects.filter(name=token.connect_method).first()
-        if not applet:
-            return rdp_options
-
-        cmdline = {
-            'app_name': applet.name,
-            'user_id': str(token.user.id),
-            'asset_id': str(token.asset.id),
-            'token_id': str(token.id)
-        }
-
-        app = '||tinker'
-        rdp_options['remoteapplicationmode:i'] = '1'
-        rdp_options['alternate shell:s'] = app
-        rdp_options['remoteapplicationprogram:s'] = app
-        rdp_options['remoteapplicationname:s'] = app
-
-        cmdline_b64 = base64.b64encode(json.dumps(cmdline).encode()).decode()
-        rdp_options['remoteapplicationcmdline:s'] = cmdline_b64
-        return rdp_options
-
     def get_rdp_file_info(self, token: ConnectionToken):
         rdp_options = {
             'full address:s': '',
@@ -114,9 +90,10 @@ class RDPFileClientProtocolURLMixin:
         rdp_options['session bpp:i'] = os.getenv('JUMPSERVER_COLOR_DEPTH', '32')
         rdp_options['audiomode:i'] = self.parse_env_bool('JUMPSERVER_DISABLE_AUDIO', 'false', '2', '0')
 
-        # 设置远程应用
-        remote_app_options = token.get_remote_app_option()
-        rdp_options.update(remote_app_options)
+        # 设置远程应用, 不是 Mstsc
+        if token.connect_method != NativeClient.mstsc:
+            remote_app_options = token.get_remote_app_option()
+            rdp_options.update(remote_app_options)
 
         # 文件名
         name = token.asset.name
@@ -160,15 +137,17 @@ class RDPFileClientProtocolURLMixin:
             'file': {}
         }
 
-        if connect_method_name == NativeClient.mstsc:
+        if connect_method_name == NativeClient.mstsc or connect_method_dict['type'] == 'applet':
             filename, content = self.get_rdp_file_info(token)
             data.update({
+                'protocol': 'rdp',
                 'file': {
                     'name': filename,
                     'content': content,
                 }
             })
         else:
+            print("Connect method: {}".format(connect_method_dict))
             endpoint = self.get_smart_endpoint(
                 protocol=connect_method_dict['endpoint_protocol'],
                 asset=token.asset
diff --git a/apps/terminal/connect_methods.py b/apps/terminal/connect_methods.py
index 70e842e26..7b17dbe4e 100644
--- a/apps/terminal/connect_methods.py
+++ b/apps/terminal/connect_methods.py
@@ -220,11 +220,10 @@ class ConnectMethodUtil:
 
         for component, component_protocol in cls.protocols().items():
             support = component_protocol['support']
+            component_web_methods = component_protocol.get('web_methods', [])
 
             for protocol in support:
                 # Web 方式
-                protocol_web_methods = set(web_methods.get(protocol, [])) \
-                                       & set(component_protocol.get('web_methods', []))
                 methods[protocol.value].extend([
                     {
                         'component': component.value,
@@ -233,7 +232,8 @@ class ConnectMethodUtil:
                         'value': method.value,
                         'label': method.label,
                     }
-                    for method in protocol_web_methods
+                    for method in web_methods.get(protocol, [])
+                    if method in component_web_methods
                 ])
 
                 # 客户端方式