perf: 修改 connect method

apps/assets/const/host.py

@@ -33,7 +33,7 @@ class HostTypes(BaseType):
     def _get_protocol_constrains(cls) -> dict:
         return {
             '*': {
-                'choices': ['ssh', 'telnet', 'vnc', 'rdp']
+                'choices': ['ssh', 'sftp', 'telnet', 'vnc', 'rdp']
             },
             cls.WINDOWS: {
                 'choices': ['rdp', 'ssh', 'vnc', 'winrm']

apps/assets/const/protocol.py

@@ -11,6 +11,7 @@ __all__ = ['Protocol']
 
 class Protocol(ChoicesMixin, models.TextChoices):
     ssh = 'ssh', 'SSH'
+    sftp = 'sftp', 'SFTP'
     rdp = 'rdp', 'RDP'
     telnet = 'telnet', 'Telnet'
     vnc = 'vnc', 'VNC'
@@ -36,17 +37,16 @@ class Protocol(ChoicesMixin, models.TextChoices):
             cls.ssh: {
                 'port': 22,
                 'secret_types': ['password', 'ssh_key'],
+            },
+            cls.sftp: {
+                'port': 22,
+                'secret_types': ['password', 'ssh_key'],
                 'setting': {
-                    'sftp_enabled': {
-                        'type': 'bool',
-                        'default': True,
-                        'label': _('SFTP enabled')
-                    },
                     'sftp_home': {
                         'type': 'str',
                         'default': '/tmp',
                         'label': _('SFTP home')
-                    },
+                    }
                 }
             },
             cls.rdp: {

apps/assets/migrations/0121_auto_20230725_1458.py

@@ -0,0 +1,77 @@
+# Generated by Django 4.1.10 on 2023-07-25 06:58
+
+from django.db import migrations
+
+
+def migrate_platforms_sftp_protocol(apps, schema_editor):
+    platform_protocol_cls = apps.get_model('assets', 'PlatformProtocol')
+    platform_cls = apps.get_model('assets', 'Platform')
+    ssh_protocols = platform_protocol_cls.objects.filter(name='ssh', setting__sftp_enabled=True)
+    platforms_has_sftp = platform_cls.objects.filter(protocols__name='sftp')
+
+    new_protocols = []
+    print("\nPlatform add sftp protocol: ")
+    for protocol in ssh_protocols:
+        protocol_setting = protocol.setting or {}
+        if protocol.platform in platforms_has_sftp:
+            continue
+
+        kwargs = {
+            'name': 'sftp',
+            'port': protocol.port,
+            'primary': False,
+            'required': False,
+            'default': True,
+            'public': True,
+            'setting': {
+                'sftp_home': protocol_setting.get('sftp_home', '/tmp'),
+            },
+            'platform': protocol.platform,
+        }
+        new_protocol = platform_protocol_cls(**kwargs)
+        new_protocols.append(new_protocol)
+        print(" - {}".format(protocol.platform.name))
+
+    new_protocols_dict = {(protocol.name, protocol.platform): protocol for protocol in new_protocols}
+    new_protocols = list(new_protocols_dict.values())
+    platform_protocol_cls.objects.bulk_create(new_protocols, ignore_conflicts=True)
+
+
+def migrate_assets_sftp_protocol(apps, schema_editor):
+    asset_cls = apps.get_model('assets', 'Asset')
+    platform_cls = apps.get_model('assets', 'Platform')
+    protocol_cls = apps.get_model('assets', 'Protocol')
+    sftp_platforms = list(platform_cls.objects.filter(protocols__name='sftp').values_list('id'))
+
+    count = 0
+    print("\nAsset add sftp protocol: ")
+    asset_ids = asset_cls.objects\
+        .filter(platform__in=sftp_platforms)\
+        .exclude(protocols__name='sftp')\
+        .distinct()\
+        .values_list('id', flat=True)
+    while True:
+        _asset_ids = asset_ids[count:count + 1000]
+        if not _asset_ids:
+            break
+        count += 1000
+
+        new_protocols = []
+        ssh_protocols = protocol_cls.objects.filter(name='ssh', asset_id__in=_asset_ids).distinct()
+        ssh_protocols_map = {protocol.asset_id: protocol for protocol in ssh_protocols}
+        for asset_id, protocol in ssh_protocols_map.items():
+            new_protocols.append(protocol_cls(name='sftp', port=protocol.port, asset_id=asset_id))
+        protocol_cls.objects.bulk_create(new_protocols, ignore_conflicts=True)
+        print(" - Add {}".format(len(new_protocols)))
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('assets', '0120_auto_20230630_1613'),
+    ]
+
+    operations = [
+        migrations.RunPython(migrate_platforms_sftp_protocol),
+        migrations.RunPython(migrate_assets_sftp_protocol),
+    ]

apps/terminal/connect_methods.py

@@ -19,14 +19,15 @@ class WebMethod(TextChoices):
     @classmethod
     def get_spec_methods(cls):
         methods = {
-            Protocol.ssh: [cls.web_cli, cls.web_sftp],
+            Protocol.sftp: [cls.web_sftp]
         }
         return methods
 
 
 class NativeClient(TextChoices):
     # Koko
-    ssh = 'ssh', 'SSH'
+    ssh = 'ssh', 'SSH CLI'
+    sftp = 'sftp', 'SFTP CLI'
     putty = 'putty', 'PuTTY'
     xshell = 'xshell', 'Xshell'
 
@@ -45,6 +46,7 @@ class NativeClient(TextChoices):
                 'default': [cls.ssh],
                 'windows': [cls.putty],
             },
+            Protocol.sftp: [cls.sftp],
             Protocol.rdp: [cls.mstsc],
             Protocol.mysql: [cls.db_client],
             Protocol.mariadb: [cls.db_client],
@@ -81,13 +83,11 @@ class NativeClient(TextChoices):
         for protocol, _clients in clients_map.items():
             if not settings.XPACK_ENABLED and protocol in xpack_protocols:
                 continue
-
             if isinstance(_clients, dict):
                 if os == 'all':
                     _clients = list(itertools.chain(*_clients.values()))
                 else:
                     _clients = _clients.get(os, _clients['default'])
-
             for client in _clients:
                 if not settings.XPACK_ENABLED and client in cls.xpack_methods():
                     continue
@@ -96,6 +96,7 @@ class NativeClient(TextChoices):
                     'label': client.label,
                     'type': 'native',
                 })
+        print("Methods: ", methods)
         return methods
 
     @classmethod
@@ -103,8 +104,10 @@ class NativeClient(TextChoices):
         username = f'JMS-{token.id}'
         commands = {
             cls.ssh: f'ssh {username}@{endpoint.host} -p {endpoint.ssh_port}',
+            cls.sftp: f'sftp {username}@{endpoint.host} -P {endpoint.ssh_port}',
             cls.putty: f'putty.exe -ssh {username}@{endpoint.host} -P {endpoint.ssh_port}',
             cls.xshell: f'xshell.exe -url ssh://{username}:{token.value}@{endpoint.host}:{endpoint.ssh_port}',
+            # 前端自己处理了
             # cls.mysql: 'mysql -h {hostname} -P {port} -u {username} -p',
             # cls.psql: {
             #     'default': 'psql -h {hostname} -p {port} -U {username} -W',
@@ -125,7 +128,6 @@ class AppletMethod:
         from .models import Applet, AppletHost
 
         methods = defaultdict(list)
-
         has_applet_hosts = AppletHost.objects.all().exists()
         applets = Applet.objects.filter(is_active=True)
         for applet in applets:
@@ -150,7 +152,7 @@ class ConnectMethodUtil:
                 'web_methods': [WebMethod.web_cli],
                 'listen': [Protocol.http, Protocol.ssh],
                 'support': [
-                    Protocol.ssh, Protocol.telnet,
+                    Protocol.ssh, Protocol.sftp, Protocol.telnet,
                     Protocol.mysql, Protocol.postgresql,
                     Protocol.sqlserver, Protocol.mariadb,
                     Protocol.redis, Protocol.mongodb,
@@ -294,7 +296,7 @@ class ConnectMethodUtil:
 
                 for listen_protocol in listen:
                     # Native method
-                    if component == TerminalType.koko and protocol.value != Protocol.ssh:
+                    if component == TerminalType.koko and protocol.value not in [Protocol.ssh, Protocol.sftp]:
                         # koko 仅支持 ssh 的 native 方式，其他数据库的 native 方式不提供
                         continue
                     methods[str(protocol)].extend([