refactor: ConnectionToken 添加 protocol 字段

apps/authentication/migrations/0013_connectiontoken_protocol.py

@@ -0,0 +1,18 @@
+# Generated by Django 3.2.14 on 2022-10-27 12:01
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('authentication', '0012_auto_20220816_1629'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='connectiontoken',
+            name='protocol',
+            field=models.CharField(choices=[('ssh', 'SSH'), ('rdp', 'RDP'), ('telnet', 'Telnet'), ('vnc', 'VNC'), ('mysql', 'MySQL'), ('mariadb', 'MariaDB'), ('oracle', 'Oracle'), ('postgresql', 'PostgreSQL'), ('sqlserver', 'SQLServer'), ('redis', 'Redis'), ('mongodb', 'MongoDB'), ('k8s', 'K8S'), ('http', 'HTTP'), ('None', ' Settings')], default='ssh', max_length=16, verbose_name='Protocol'),
+        ),
+    ]

apps/authentication/models/connection_token.py

@@ -9,6 +9,7 @@ from django.db import models
 from common.utils import lazyproperty
 from common.utils.timezone import as_current_tz
 from common.db.models import JMSBaseModel
+from assets.const import Protocol
 
 
 def date_expired_default():
@@ -26,10 +27,14 @@ class ConnectionToken(OrgModelMixin, JMSBaseModel):
     )
     user_display = models.CharField(max_length=128, default='', verbose_name=_("User display"))
     asset_display = models.CharField(max_length=128, default='', verbose_name=_("Asset display"))
-    protocol = ''
     account = models.CharField(max_length=128, default='', verbose_name=_("Account"))
+    protocol = models.CharField(
+        choices=Protocol.choices, max_length=16, default=Protocol.ssh, verbose_name=_("Protocol")
+    )
     secret = models.CharField(max_length=64, default='', verbose_name=_("Secret"))
-    date_expired = models.DateTimeField(default=date_expired_default, verbose_name=_("Date expired"))
+    date_expired = models.DateTimeField(
+        default=date_expired_default, verbose_name=_("Date expired")
+    )
 
     class Meta:
         ordering = ('-date_expired',)

apps/authentication/serializers/connection_token.py

@@ -17,7 +17,6 @@ __all__ = [
 
 
 class ConnectionTokenSerializer(OrgResourceModelSerializerMixin):
-    type_display = serializers.ReadOnlyField(source='get_type_display', label=_("Type display"))
     is_valid = serializers.BooleanField(read_only=True, label=_('Validity'))
     expire_time = serializers.IntegerField(read_only=True, label=_('Expired time'))
 
@@ -29,13 +28,13 @@ class ConnectionTokenSerializer(OrgResourceModelSerializerMixin):
             'created_by', 'updated_by', 'org_id', 'org_name',
         ]
         fields_fk = [
-            'user', 'system_user', 'asset', 'application',
+            'user', 'system_user', 'asset',
         ]
         read_only_fields = [
             # 普通 Token 不支持指定 user
             'user', 'is_valid', 'expire_time',
-            'type_display', 'user_display', 'system_user_display',
+            'user_display', 'system_user_display',
-            'asset_display', 'application_display',
+            'asset_display',
         ]
         fields = fields_small + fields_fk + read_only_fields
 
@@ -54,28 +53,23 @@ class ConnectionTokenSerializer(OrgResourceModelSerializerMixin):
         return self.request_user
 
     def construct_internal_fields_attrs(self, attrs):
-        user = self.get_user(attrs)
-        system_user = attrs.get('system_user') or ''
         asset = attrs.get('asset') or ''
-        application = attrs.get('application') or ''
+        asset_display = pretty_string(str(asset), max_length=128)
+        user = self.get_user(attrs)
+        user_display = pretty_string(str(user), max_length=128)
         secret = attrs.get('secret') or random_string(16)
         date_expired = attrs.get('date_expired') or ConnectionToken.get_default_date_expired()
-
-        if isinstance(asset, Asset):
-            tp = ConnectionToken.Type.asset
         org_id = asset.org_id
-        else:
+        if not isinstance(asset, Asset):
-            raise serializers.ValidationError(_('Asset or application required'))
+            error = ''
+            raise serializers.ValidationError(error)
 
         return {
-            'type': tp,
             'user': user,
             'secret': secret,
+            'user_display': user_display,
+            'asset_display': asset_display,
             'date_expired': date_expired,
-            'user_display': pretty_string(str(user), max_length=128),
-            'system_user_display': pretty_string(str(system_user), max_length=128),
-            'asset_display': pretty_string(str(asset), max_length=128),
-            'application_display': pretty_string(str(application), max_length=128),
             'org_id': org_id,
         }
 
@@ -155,7 +149,6 @@ class ConnectionTokenSecretSerializer(OrgResourceModelSerializerMixin):
     user = ConnectionTokenUserSerializer(read_only=True)
     asset = ConnectionTokenAssetSerializer(read_only=True)
     remote_app = ConnectionTokenRemoteAppSerializer(read_only=True)
-    account = serializers.CharField(read_only=True)
     gateway = ConnectionTokenGatewaySerializer(read_only=True)
     domain = ConnectionTokenDomainSerializer(read_only=True)
     cmd_filter_rules = ConnectionTokenCmdFilterRuleSerializer(many=True)
@@ -165,6 +158,6 @@ class ConnectionTokenSecretSerializer(OrgResourceModelSerializerMixin):
     class Meta:
         model = ConnectionToken
         fields = [
-            'id', 'secret', 'type', 'user', 'asset', 'account',
+            'id', 'secret', 'type', 'user', 'asset', 'account', 'protocol',
             'cmd_filter_rules', 'domain', 'gateway', 'actions', 'expired_at',
         ]

apps/perms/utils/account.py

@@ -51,7 +51,6 @@ class PermAccountUtil(AssetPermissionUtil):
             user, asset, with_actions=True, with_perms=True
         )
         perm = perms.first()
-        # Todo: 后面可能需要加上 protocol 进行过滤, 因为同名的账号协议是不一样可能会存在多个
         account = accounts.filter(username=account_username).first()
         actions = account.actions if account else []
         expire_at = perm.date_expired if perm else time.time()