Merge pull request #11091 from jumpserver/pr@dev@ssh_to_sftp

perf: 修改 sftp 协议

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,12 +46,12 @@ 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],
             Protocol.redis: [cls.db_client],
             Protocol.mongodb: [cls.db_client],
-
             Protocol.oracle: [cls.db_client],
             Protocol.postgresql: [cls.db_client],
         }
@@ -81,13 +82,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
@@ -103,8 +102,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 +126,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:
@@ -148,12 +148,18 @@ class ConnectMethodUtil:
         protocols = {
             TerminalType.koko: {
                 'web_methods': [WebMethod.web_cli],
-                'listen': [Protocol.http, Protocol.ssh],
+                'listen': [Protocol.http, Protocol.ssh, Protocol.sftp],
                 'support': [
-                    Protocol.ssh, Protocol.telnet,
+                    Protocol.ssh, Protocol.telnet, Protocol.sftp,
                     Protocol.redis, Protocol.mongodb,
                     Protocol.k8s, Protocol.clickhouse,
                 ],
+                # 限制客户端的协议，比如 koko 虽然也支持 数据库的 ssh 连接，但是不再这里拉起
+                # Listen协议: [Asset协议]
+                'client_limits': {
+                    Protocol.sftp: [Protocol.sftp],
+                    Protocol.ssh: [Protocol.ssh, Protocol.telnet],
+                },
                 'match': 'm2m'
             },
             TerminalType.chen: {
@@ -260,20 +266,20 @@ class ConnectMethodUtil:
 
         methods = defaultdict(list)
         spec_web_methods = WebMethod.get_spec_methods()
-        native_methods = NativeClient.get_methods(os)
         applet_methods = AppletMethod.get_methods()
+        native_methods = NativeClient.get_methods(os=os)
 
         for component, component_protocol in cls.components().items():
             support = component_protocol['support']
-            component_web_methods = component_protocol.get('web_methods', [])
+            default_web_methods = component_protocol.get('web_methods', [])
+            client_limits = component_protocol.get('client_limits', {})
 
-            for protocol in support:
+            for asset_protocol in support:
                 # Web 方式
-                web_methods = spec_web_methods.get(protocol, None)
+                web_methods = spec_web_methods.get(asset_protocol, [])
-                if web_methods is None:
+                if not web_methods:
-                    web_methods = component_web_methods
+                    web_methods = default_web_methods
-
+                methods[str(asset_protocol)].extend([
-                methods[str(protocol)].extend([
                     {
                         'component': component.value,
                         'type': 'web',
@@ -286,31 +292,32 @@ class ConnectMethodUtil:
 
                 # 客户端方式
                 if component_protocol['match'] == 'map':
-                    listen = [protocol]
+                    listen = [asset_protocol]
                 else:
                     listen = component_protocol['listen']
 
                 for listen_protocol in listen:
-                    # Native method
+                    limits = client_limits.get(listen_protocol, [])
-                    if component == TerminalType.koko and protocol.value != Protocol.ssh:
+                    if limits and asset_protocol not in limits:
-                        # koko 仅支持 ssh 的 native 方式，其他数据库的 native 方式不提供
                         continue
-                    methods[str(protocol)].extend([
+                    # Native method
+                    client_methods = native_methods.get(listen_protocol, [])
+                    methods[str(asset_protocol)].extend([
                         {
                             'component': component.value,
                             'type': 'native',
                             'endpoint_protocol': listen_protocol,
                             **method
                         }
-                        for method in native_methods[listen_protocol]
+                        for method in client_methods
                     ])
 
         # 远程应用方式，这个只有 tinker 提供，并且协议可能是自定义的
-        for protocol, applet_methods in applet_methods.items():
+        for asset_protocol, applet_methods in applet_methods.items():
             for method in applet_methods:
                 method['listen'] = 'rdp'
                 method['component'] = TerminalType.tinker.value
-            methods[protocol].extend(applet_methods)
+            methods[asset_protocol].extend(applet_methods)
 
         cls._all_methods[os] = methods
         return methods