From c135837372a40b29fe7fad70105dc7089de1b9ea Mon Sep 17 00:00:00 2001 From: ibuler Date: Tue, 25 Jul 2023 17:12:06 +0800 Subject: [PATCH 1/2] =?UTF-8?q?perf:=20=E4=BF=AE=E6=94=B9=20connect=20meth?= =?UTF-8?q?od?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/assets/const/host.py | 2 +- apps/assets/const/protocol.py | 12 +-- .../migrations/0121_auto_20230725_1458.py | 77 +++++++++++++++++++ apps/terminal/connect_methods.py | 16 ++-- 4 files changed, 93 insertions(+), 14 deletions(-) create mode 100644 apps/assets/migrations/0121_auto_20230725_1458.py diff --git a/apps/assets/const/host.py b/apps/assets/const/host.py index afb92a447..cb841bbc2 100644 --- a/apps/assets/const/host.py +++ b/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'] diff --git a/apps/assets/const/protocol.py b/apps/assets/const/protocol.py index aface581c..f405decb8 100644 --- a/apps/assets/const/protocol.py +++ b/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: { diff --git a/apps/assets/migrations/0121_auto_20230725_1458.py b/apps/assets/migrations/0121_auto_20230725_1458.py new file mode 100644 index 000000000..d32eb86a1 --- /dev/null +++ b/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), + ] diff --git a/apps/terminal/connect_methods.py b/apps/terminal/connect_methods.py index ed49d722e..1673c55f3 100644 --- a/apps/terminal/connect_methods.py +++ b/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([ From b0b6d19bc0125652bf10f64163863604c12f142e Mon Sep 17 00:00:00 2001 From: ibuler Date: Wed, 26 Jul 2023 15:31:02 +0800 Subject: [PATCH 2/2] =?UTF-8?q?perf:=20=E4=BF=AE=E6=94=B9=20sftp=20?= =?UTF-8?q?=E5=8D=8F=E8=AE=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/terminal/connect_methods.py | 45 ++++++++++++++++++-------------- 1 file changed, 25 insertions(+), 20 deletions(-) diff --git a/apps/terminal/connect_methods.py b/apps/terminal/connect_methods.py index bc8ffbc36..aab6f6c75 100644 --- a/apps/terminal/connect_methods.py +++ b/apps/terminal/connect_methods.py @@ -52,7 +52,6 @@ class NativeClient(TextChoices): 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], } @@ -96,7 +95,6 @@ class NativeClient(TextChoices): 'label': client.label, 'type': 'native', }) - print("Methods: ", methods) return methods @classmethod @@ -150,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.sftp, 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: { @@ -262,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) - if web_methods is None: - web_methods = component_web_methods - - methods[str(protocol)].extend([ + web_methods = spec_web_methods.get(asset_protocol, []) + if not web_methods: + web_methods = default_web_methods + methods[str(asset_protocol)].extend([ { 'component': component.value, 'type': 'web', @@ -288,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 - if component == TerminalType.koko and protocol.value not in [Protocol.ssh, Protocol.sftp]: - # koko 仅支持 ssh 的 native 方式，其他数据库的 native 方式不提供 + limits = client_limits.get(listen_protocol, []) + if limits and asset_protocol not in limits: 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