From fb653f93db92bccf1d860a7b65995dffc4afe088 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Wed, 16 Nov 2022 21:05:15 +0800
Subject: [PATCH] =?UTF-8?q?pref:=20=E4=BF=AE=E6=94=B9=E6=B7=BB=E5=8A=A0=20?=
 =?UTF-8?q?connect=20methods?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/terminal/api/component/endpoint.py | 36 ++++++++++------
 apps/terminal/const.py                  | 55 ++++++++++++++++++++++++-
 apps/terminal/serializers/terminal.py   |  8 +++-
 3 files changed, 83 insertions(+), 16 deletions(-)

diff --git a/apps/terminal/api/component/endpoint.py b/apps/terminal/api/component/endpoint.py
index 364cd803e..c81cbc1c3 100644
--- a/apps/terminal/api/component/endpoint.py
+++ b/apps/terminal/api/component/endpoint.py
@@ -1,17 +1,18 @@
-from rest_framework.decorators import action
-from rest_framework.response import Response
-from rest_framework import status
-from rest_framework.request import Request
-
-from common.drf.api import JMSBulkModelViewSet
-from common.permissions import IsValidUserOrConnectionToken
-from django.utils.translation import ugettext_lazy as _
 from django.shortcuts import get_object_or_404
-from assets.models import Asset
-from orgs.utils import tmp_to_root_org
-from terminal.models import Session, Endpoint, EndpointRule
-from terminal import serializers
+from django.utils.translation import ugettext_lazy as _
+from rest_framework import status
+from rest_framework.decorators import action
+from rest_framework.generics import ListAPIView
+from rest_framework.request import Request
+from rest_framework.response import Response
 
+from assets.models import Asset
+from common.drf.api import JMSBulkModelViewSet
+from common.permissions import IsValidUser
+from common.permissions import IsValidUserOrConnectionToken
+from orgs.utils import tmp_to_root_org
+from terminal import serializers
+from terminal.models import Session, Endpoint, EndpointRule
 
 __all__ = ['EndpointViewSet', 'EndpointRuleViewSet']
 
@@ -97,3 +98,14 @@ class EndpointRuleViewSet(JMSBulkModelViewSet):
     search_fields = filterset_fields
     serializer_class = serializers.EndpointRuleSerializer
     queryset = EndpointRule.objects.all()
+
+
+class ConnectMethodListApi(ListAPIView):
+    permission_classes = (IsValidUser,)
+    serializer_class = serializers.ProtocolConnectMethodsSerializer
+
+    def get_queryset(self):
+        protocol = self.request.query_params.get('protocol')
+        if not protocol:
+            return []
+        return Protocol.objects.filter(name=protocol)
diff --git a/apps/terminal/const.py b/apps/terminal/const.py
index acea15238..d44303f9f 100644
--- a/apps/terminal/const.py
+++ b/apps/terminal/const.py
@@ -4,6 +4,9 @@
 from django.db.models import TextChoices
 from django.utils.translation import ugettext_lazy as _
 
+from assets.const import Protocol
+
+
 # Replay & Command Storage Choices
 # --------------------------------
 
@@ -48,11 +51,59 @@ class TerminalType(TextChoices):
     lion = 'lion', 'Lion'
     core = 'core', 'Core'
     celery = 'celery', 'Celery'
-    magnus = 'magnus',  'Magnus'
-    razor = 'razor',  'Razor'
+    magnus = 'magnus', 'Magnus'
+    razor = 'razor', 'Razor'
     tinker = 'tinker', 'Tinker'
 
     @classmethod
     def types(cls):
         return set(dict(cls.choices).keys())
 
+
+class NativeClient:
+    # Koko
+    ssh = 'ssh', 'ssh'
+    putty = 'putty', 'PuTTY'
+    xshell = 'xshell', 'Xshell'
+
+    # Magnus
+    mysql = 'mysql', 'MySQL Client'
+    psql = 'psql', 'psql'
+    sqlplus = 'sqlplus', 'sqlplus'
+    redis = 'redis-cli', 'redis-cli'
+
+    # Razor
+    mstsc = 'mstsc', 'Remote Desktop'
+
+    @classmethod
+    def commands(cls, name, os):
+        return {
+            'ssh': 'ssh {username}@{hostname} -p {port}',
+            'putty': 'putty -ssh {username}@{hostname} -P {port}',
+            'xshell': '-url ssh://root:passwd@192.168.10.100',
+            'mysql': 'mysql -h {hostname} -P {port} -u {username} -p',
+            'psql': {
+                'default': 'psql -h {hostname} -p {port} -U {username} -W',
+                'windows': 'psql /h {hostname} /p {port} /U {username} -W',
+            },
+            'sqlplus': 'sqlplus {username}/{password}@{hostname}:{port}',
+            'redis': 'redis-cli -h {hostname} -p {port} -a {password}',
+            'mstsc': 'mstsc /v:{hostname}:{port}',
+        }
+
+
+class ConnectMethod(TextChoices):
+    web_cli = 'web_cli', _('Web CLI')
+    web_gui = 'web_gui', _('Web GUI')
+    native_client = 'native_client', _('Native Client')
+
+    @classmethod
+    def methods(cls):
+        return {
+            Protocol.ssh: [cls.web_cli, cls.native_client],
+            Protocol.rdp: ([cls.web_gui], [cls.native_client]),
+            Protocol.vnc: [cls.web_gui],
+            Protocol.telnet: [cls.web_cli, cls.native_client],
+            Protocol.mysql: [cls.web_cli, cls.web_gui, cls.native_client],
+            Protocol.sqlserver: [cls.web_cli, cls.web_gui],
+        }
diff --git a/apps/terminal/serializers/terminal.py b/apps/terminal/serializers/terminal.py
index 4b2e3614e..f12990618 100644
--- a/apps/terminal/serializers/terminal.py
+++ b/apps/terminal/serializers/terminal.py
@@ -1,8 +1,8 @@
-from rest_framework import serializers
 from django.utils.translation import ugettext_lazy as _
+from rest_framework import serializers
 
-from common.drf.serializers import BulkModelSerializer
 from common.drf.fields import LabeledChoiceField
+from common.drf.serializers import BulkModelSerializer
 from common.utils import get_request_ip, pretty_string, is_uuid
 from users.serializers import ServiceAccountSerializer
 from .. import const
@@ -133,3 +133,7 @@ class TerminalRegistrationSerializer(serializers.ModelSerializer):
         instance.replay_storage = ReplayStorage.default().name
         instance.save()
         return instance
+
+
+class ConnectMethodSerializer(serializers.Serializer):
+    name = serializers.CharField(max_length=128)