refactor: 修改 ConnectionToken 关联的逻辑(1)

apps/authentication/api/connection_token.py

@@ -1,6 +1,7 @@
 import abc
 import os
 import json
+import time
 import base64
 import urllib.parse
 from django.http import HttpResponse
@@ -41,17 +42,25 @@ class ConnectionTokenMixin:
     def get_request_resources(self, serializer):
         user = self.get_request_resource_user(serializer)
         asset = serializer.validated_data.get('asset')
-        application = serializer.validated_data.get('application')
-        system_user = serializer.validated_data.get('system_user')
-        return user, asset, application, system_user
+        account = serializer.validated_data.get('account')
+        return user, asset, account
 
     @staticmethod
-    def check_user_has_resource_permission(user, asset, application, system_user):
-        from perms.utils.asset import has_asset_system_permission
+    def check_user_has_resource_permission(user, asset, account):
+        from perms.utils.account import PermAccountUtil
+        if not asset or not user:
+            error = ''
+            raise PermissionDenied(error)
 
-        if asset and not has_asset_system_permission(user, asset, system_user):
-            error = f'User not has this asset and system user permission: ' \
-                    f'user={user.id} system_user={system_user.id} asset={asset.id}'
+        actions, expire_at = PermAccountUtil().validate_permission(
+            user, asset, account_username=account
+        )
+        if not actions:
+            error = ''
+            raise PermissionDenied(error)
+
+        if expire_at < time.time():
+            error = ''
             raise PermissionDenied(error)
 
     def get_smart_endpoint(self, protocol, asset=None, application=None):
@@ -69,13 +78,12 @@ class ConnectionTokenMixin:
         return true_value if is_true(os.getenv(env_key, env_default)) else false_value
 
     def get_client_protocol_data(self, token: ConnectionToken):
-        from assets.models import SystemUser
-        protocol = token.system_user.protocol
+        protocol = token.protocol
         username = token.user.username
         rdp_config = ssh_token = ''
-        if protocol == SystemUser.Protocol.rdp:
+        if protocol == 'rdp':
             filename, rdp_config = self.get_rdp_file_info(token)
-        elif protocol == SystemUser.Protocol.ssh:
+        elif protocol == 'ssh':
             filename, ssh_token = self.get_ssh_token(token)
         else:
             raise ValueError('Protocol not support: {}'.format(protocol))
@@ -134,15 +142,12 @@ class ConnectionTokenMixin:
         rdp_options['screen mode id:i'] = '2' if full_screen else '1'
 
         # 设置 RDP Server 地址
-        endpoint = self.get_smart_endpoint(
-            protocol='rdp', asset=token.asset, application=token.application
-        )
+        endpoint = self.get_smart_endpoint(protocol='rdp', asset=token.asset)
         rdp_options['full address:s'] = f'{endpoint.host}:{endpoint.rdp_port}'
 
         # 设置用户名
         rdp_options['username:s'] = '{}|{}'.format(token.user.username, str(token.id))
-        if token.system_user.ad_domain:
-            rdp_options['domain:s'] = token.system_user.ad_domain
+        # rdp_options['domain:s'] = token.account_ad_domain
 
         # 设置宽高
         height = self.request.query_params.get('height')
@@ -158,13 +163,12 @@ class ConnectionTokenMixin:
 
         if token.asset:
             name = token.asset.name
-        elif token.application and token.application.category_remote_app:
-            app = '||jmservisor'
-            name = token.application.name
-            rdp_options['remoteapplicationmode:i'] = '1'
-            rdp_options['alternate shell:s'] = app
-            rdp_options['remoteapplicationprogram:s'] = app
-            rdp_options['remoteapplicationname:s'] = name
+            # remote-app
+            # app = '||jmservisor'
+            # rdp_options['remoteapplicationmode:i'] = '1'
+            # rdp_options['alternate shell:s'] = app
+            # rdp_options['remoteapplicationprogram:s'] = app
+            # rdp_options['remoteapplicationname:s'] = name
         else:
             name = '*'
         prefix_name = f'{token.user.username}-{name}'
@@ -188,16 +192,12 @@ class ConnectionTokenMixin:
     def get_ssh_token(self, token: ConnectionToken):
         if token.asset:
             name = token.asset.name
-        elif token.application:
-            name = token.application.name
         else:
             name = '*'
         prefix_name = f'{token.user.username}-{name}'
         filename = self.get_connect_filename(prefix_name)
 
-        endpoint = self.get_smart_endpoint(
-            protocol='ssh', asset=token.asset, application=token.application
-        )
+        endpoint = self.get_smart_endpoint(protocol='ssh', asset=token.asset)
         data = {
             'ip': endpoint.host,
             'port': str(endpoint.ssh_port),
@@ -251,8 +251,8 @@ class ConnectionTokenViewSet(ConnectionTokenMixin, RootOrgViewMixin, JMSModelVie
         return token
 
     def perform_create(self, serializer):
-        user, asset, application, system_user = self.get_request_resources(serializer)
-        self.check_user_has_resource_permission(user, asset, application, system_user)
+        user, asset, account = self.get_request_resources(serializer)
+        self.check_user_has_resource_permission(user, asset, account)
         return super(ConnectionTokenViewSet, self).perform_create(serializer)
 
     @action(methods=['POST'], detail=False, url_path='secret-info/detail')
@@ -264,7 +264,6 @@ class ConnectionTokenViewSet(ConnectionTokenMixin, RootOrgViewMixin, JMSModelVie
         token_id = request.data.get('token') or ''
         token = get_object_or_404(ConnectionToken, pk=token_id)
         self.check_token_valid(token)
-        token.load_system_user_auth()
         serializer = self.get_serializer(instance=token)
         return Response(serializer.data, status=status.HTTP_200_OK)
 

apps/authentication/models.py

@@ -1,3 +1,4 @@
+import time
 import uuid
 from datetime import datetime, timedelta
 from django.utils import timezone
@@ -63,22 +64,20 @@ def date_expired_default():
 
 
 class ConnectionToken(OrgModelMixin, JMSBaseModel):
-    secret = models.CharField(max_length=64, default='', verbose_name=_("Secret"))
-    date_expired = models.DateTimeField(
-        default=date_expired_default, verbose_name=_("Date expired")
-    )
-
     user = models.ForeignKey(
-        'users.User', on_delete=models.SET_NULL, verbose_name=_('User'),
-        related_name='connection_tokens', null=True, blank=True
+        'users.User', on_delete=models.SET_NULL,  null=True, blank=True,
+        related_name='connection_tokens', verbose_name=_('User')
+    )
+    asset = models.ForeignKey(
+        'assets.Asset', on_delete=models.SET_NULL, null=True, blank=True,
+        related_name='connection_tokens', verbose_name=_('Asset'),
     )
     user_display = models.CharField(max_length=128, default='', verbose_name=_("User display"))
-    asset = models.ForeignKey(
-        'assets.Asset', on_delete=models.SET_NULL, verbose_name=_('Asset'),
-        related_name='connection_tokens', null=True, blank=True
-    )
     asset_display = models.CharField(max_length=128, default='', verbose_name=_("Asset display"))
+    protocol = ''
     account = models.CharField(max_length=128, default='', verbose_name=_("Account"))
+    secret = models.CharField(max_length=64, default='', verbose_name=_("Secret"))
+    date_expired = models.DateTimeField(default=date_expired_default, verbose_name=_("Date expired"))
 
     class Meta:
         ordering = ('-date_expired',)
@@ -87,10 +86,6 @@ class ConnectionToken(OrgModelMixin, JMSBaseModel):
             ('view_connectiontokensecret', _('Can view connection token secret'))
         ]
 
-    @classmethod
-    def get_default_date_expired(cls):
-        return date_expired_default()
-
     @property
     def is_expired(self):
         return self.date_expired < timezone.now()
@@ -103,32 +98,32 @@ class ConnectionToken(OrgModelMixin, JMSBaseModel):
             seconds = 0
         return int(seconds)
 
-    def expire(self):
-        self.date_expired = timezone.now()
-        self.save()
-
     @property
     def is_valid(self):
         return not self.is_expired
 
-    def is_type(self, tp):
-        return self.type == tp
+    @classmethod
+    def get_default_date_expired(cls):
+        return date_expired_default()
+
+    def expire(self):
+        self.date_expired = timezone.now()
+        self.save()
 
     def renewal(self):
         """ 续期 Token，将来支持用户自定义创建 token 后，续期策略要修改 """
         self.date_expired = self.get_default_date_expired()
         self.save()
 
-    actions = expired_at = None  # actions 和 expired_at 在 check_valid() 中赋值
+    # actions 和 expired_at 在 check_valid() 中赋值
+    actions = expire_at = None
 
     def check_valid(self):
-        from perms.utils.permission import validate_permission as asset_validate_permission
-
+        from perms.utils.account import PermAccountUtil
         if self.is_expired:
             is_valid = False
             error = _('Connection token expired at: {}').format(as_current_tz(self.date_expired))
             return is_valid, error
-
         if not self.user:
             is_valid = False
             error = _('User not exists')
@@ -137,44 +132,33 @@ class ConnectionToken(OrgModelMixin, JMSBaseModel):
             is_valid = False
             error = _('User invalid, disabled or expired')
             return is_valid, error
-
+        if not self.asset:
+            is_valid = False
+            error = _('Asset not exists')
+            return is_valid, error
+        if not self.asset.is_active:
+            is_valid = False
+            error = _('Asset inactive')
+            return is_valid, error
         if not self.account:
             is_valid = False
             error = _('Account not exists')
             return is_valid, error
 
-        if not self.asset:
-            is_valid = False
-            error = _('Asset not exists')
-            return is_valid, error
-
-        if not self.asset.is_active:
-            is_valid = False
-            error = _('Asset inactive')
-            return is_valid, error
-
-        has_perm, actions, expired_at = asset_validate_permission(
+        actions, expire_at = PermAccountUtil().validate_permission(
             self.user, self.asset, self.account
         )
-        if not has_perm:
+        if not actions or expire_at < time.time():
             is_valid = False
             error = _('User has no permission to access asset or permission expired')
             return is_valid, error
         self.actions = actions
-        self.expired_at = expired_at
+        self.expire_at = expire_at
         return True, ''
 
     @lazyproperty
     def domain(self):
-        if self.asset:
-            return self.asset.domain
-        if not self.application:
-            return
-        if self.application.category_remote_app:
-            asset = self.application.get_remote_app_asset()
-            domain = asset.domain if asset else None
-        else:
-            domain = self.application.domain
+        domain = self.asset.domain if self.asset else None
         return domain
 
     @lazyproperty
@@ -185,41 +169,18 @@ class ConnectionToken(OrgModelMixin, JMSBaseModel):
         self.domain: Domain
         return self.domain.random_gateway()
 
-    @lazyproperty
-    def remote_app(self):
-        if not self.application:
-            return {}
-        if not self.application.category_remote_app:
-            return {}
-        return self.application.get_rdp_remote_app_setting()
-
-    @lazyproperty
-    def asset_or_remote_app_asset(self):
-        if self.asset:
-            return self.asset
-        if self.application and self.application.category_remote_app:
-            return self.application.get_remote_app_asset()
-
     @lazyproperty
     def cmd_filter_rules(self):
         from assets.models import CommandFilterRule
         kwargs = {
             'user_id': self.user.id,
-            'system_user_id': self.system_user.id,
+            'account': self.account,
         }
         if self.asset:
             kwargs['asset_id'] = self.asset.id
-        elif self.application:
-            kwargs['application_id'] = self.application_id
         rules = CommandFilterRule.get_queryset(**kwargs)
         return rules
 
-    def load_system_user_auth(self):
-        if self.asset:
-            self.system_user.load_asset_more_auth(self.asset.id, self.user.username, self.user.id)
-        elif self.application:
-            self.system_user.load_app_more_auth(self.application.id, self.user.username, self.user.id)
-
 
 class TempToken(JMSBaseModel):
     username = models.CharField(max_length=128, verbose_name=_("Username"))

apps/authentication/serializers/connection_token.py

@@ -165,6 +165,6 @@ class ConnectionTokenSecretSerializer(OrgResourceModelSerializerMixin):
     class Meta:
         model = ConnectionToken
         fields = [
-            'id', 'secret', 'type', 'user', 'asset', 'application', 'system_user',
-            'remote_app', 'cmd_filter_rules', 'domain', 'gateway', 'actions', 'expired_at',
+            'id', 'secret', 'type', 'user', 'asset', 'account',
+            'cmd_filter_rules', 'domain', 'gateway', 'actions', 'expired_at',
         ]