perf: 修改命令过滤相关的Model, CommandFilterACL, CommandGroup; 修改Model QuerySet 相关的方法;

apps/acls/api/login_asset_check.py

@@ -31,13 +31,13 @@ class LoginAssetCheckAPI(CreateAPIView):
 
     def check_confirm(self):
         with tmp_to_org(self.serializer.asset.org):
-            acl = LoginAssetACL.objects \
-                .filter(action=LoginAssetACL.ActionChoices.review) \
-                .filter_user(self.serializer.user) \
-                .filter_asset(self.serializer.asset) \
-                .filter_account(self.serializer.validated_data.get('account_username')) \
-                .valid() \
-                .first()
+            kwargs = {
+                'user': self.serializer.user,
+                'asset': self.serializer.asset,
+                'account_username': self.serializer.validated_data.get('account_username'),
+                'action': LoginAssetACL.ActionChoices.review
+            }
+            acl = LoginAssetACL.filter_queryset(**kwargs).valid().first()
         if acl:
             need_confirm = True
             response_data = self._get_response_data_of_need_confirm(acl)

apps/acls/migrations/0009_auto_20221204_0001.py

@@ -0,0 +1,22 @@
+# Generated by Django 3.2.14 on 2022-12-03 16:01
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('acls', '0008_commandgroup_comment'),
+    ]
+
+    operations = [
+        migrations.AlterModelOptions(
+            name='commandgroup',
+            options={'verbose_name': 'Command group'},
+        ),
+        migrations.RenameField(
+            model_name='commandfilteracl',
+            old_name='commands',
+            new_name='command_groups',
+        ),
+    ]

apps/acls/models/base.py

@@ -5,8 +5,15 @@ from django.utils.translation import ugettext_lazy as _
 
 from common.mixins import CommonModelMixin
 from common.utils import contains_ip
+from orgs.mixins.models import OrgModelMixin
 
-__all__ = ['BaseACL', 'BaseACLQuerySet', 'ACLManager', 'AssetAccountUserACLQuerySet']
+__all__ = [
+    'ACLManager',
+    'BaseACL',
+    'BaseACLQuerySet',
+    'UserAssetAccountBaseACL',
+    'UserAssetAccountACLQuerySet'
+]
 
 
 class ActionChoices(models.TextChoices):
@@ -29,30 +36,35 @@ class BaseACLQuerySet(models.QuerySet):
         return self.inactive()
 
 
-class AssetAccountUserACLQuerySet(BaseACLQuerySet):
-    def filter_user(self, user):
-        return self.filter(
-            Q(users__username_group__contains=user.username) |
+class UserAssetAccountACLQuerySet(BaseACLQuerySet):
+    def filter_user(self, username):
+        q = Q(users__username_group__contains=username) | \
             Q(users__username_group__contains='*')
-        )
-
-    def filter_asset(self, asset):
-        queryset = self.filter(
-            Q(assets__name_group__contains=asset.name) |
-            Q(assets__name_group__contains='*')
-        )
-        ids = [
-            q.id for q in queryset
-            if contains_ip(asset.address, q.assets.get('address_group', []))
-        ]
-        queryset = self.filter(id__in=ids)
+        return self.filter(q)
+
+    def filter_asset(self, name=None, address=None):
+        queryset = self.filter()
+        if name:
+            q = Q(assets__name_group__contains=name) | \
+                Q(assets__name_group__contains='*')
+            queryset = queryset.filter(q)
+        if address:
+            ids = [
+                q.id for q in queryset
+                if contains_ip(address, q.assets.get('address_group', []))
+            ]
+            queryset = queryset.filter(id__in=ids)
         return queryset
 
-    def filter_account(self, account_username):
-        return self.filter(
-            Q(accounts__username_group__contains=account_username) |
-            Q(accounts__username_group__contains='*')
-        )
+    def filter_account(self, name=None, username=None):
+        q = Q()
+        if name:
+            q &= Q(accounts__name_group__contains=name) | \
+                 Q(accounts__name_group__contains='*')
+        if username:
+            q &= Q(accounts__username_group__contains=username) | \
+                 Q(accounts__username_group__contains='*')
+        return self.filter(q)
 
 
 class ACLManager(models.Manager):
@@ -72,8 +84,43 @@ class BaseACL(CommonModelMixin):
     is_active = models.BooleanField(default=True, verbose_name=_("Active"))
     comment = models.TextField(default='', blank=True, verbose_name=_('Comment'))
 
-    objects = ACLManager.from_queryset(BaseACLQuerySet)()
     ActionChoices = ActionChoices
+    objects = ACLManager.from_queryset(BaseACLQuerySet)()
 
     class Meta:
         abstract = True
+
+
+class UserAssetAccountBaseACL(BaseACL, OrgModelMixin):
+    # username_group
+    users = models.JSONField(verbose_name=_('User'))
+    # name_group, address_group
+    assets = models.JSONField(verbose_name=_('Asset'))
+    # name_group, username_group
+    accounts = models.JSONField(verbose_name=_('Account'))
+
+    objects = ACLManager.from_queryset(UserAssetAccountACLQuerySet)()
+
+    class Meta:
+        abstract = True
+
+    @classmethod
+    def filter_queryset(cls, user=None, asset=None, account=None, account_username=None, **kwargs):
+        queryset = cls.objects.all()
+        org_id = None
+        if user:
+            queryset = queryset.filter_user(user.username)
+        if asset:
+            org_id = asset.org_id
+            queryset = queryset.filter_asset(asset.name, asset.address)
+        if account:
+            org_id = account.org_id
+            queryset = queryset.filter_account(account.name, account.username)
+        if account_username:
+            queryset = queryset.filter_account(username=account_username)
+        if org_id:
+            kwargs['org_id'] = org_id
+        if kwargs:
+            queryset = queryset.filter(**kwargs)
+        return queryset
+

apps/acls/models/command_acl.py

@@ -3,37 +3,41 @@
 import re
 
 from django.db import models
-from django.db.models import Q
 from django.utils.translation import ugettext_lazy as _
 
-from common.utils import lazyproperty, get_logger, get_object_or_none
+from common.utils import lazyproperty, get_logger
 from orgs.mixins.models import JMSOrgBaseModel
-from orgs.mixins.models import OrgModelMixin
-from users.models import User, UserGroup
-from .base import BaseACL, AssetAccountUserACLQuerySet, ACLManager
+
+from .base import UserAssetAccountBaseACL, UserAssetAccountACLQuerySet, ACLManager
 
 logger = get_logger(__file__)
 
 
-class CommandGroup(JMSOrgBaseModel):
-    class Type(models.TextChoices):
-        command = 'command', _('Command')
-        regex = 'regex', _('Regex')
+class TypeChoices(models.TextChoices):
+    command = 'command', _('Command')
+    regex = 'regex', _('Regex')
+
 
+class CommandGroup(JMSOrgBaseModel):
     name = models.CharField(max_length=128, verbose_name=_("Name"))
-    type = models.CharField(max_length=16, default=Type.command, choices=Type.choices, verbose_name=_("Type"))
+    type = models.CharField(
+        max_length=16, default=TypeChoices.command, choices=TypeChoices.choices,
+        verbose_name=_("Type")
+    )
     content = models.TextField(verbose_name=_("Content"), help_text=_("One line one command"))
     ignore_case = models.BooleanField(default=True, verbose_name=_('Ignore case'))
     comment = models.TextField(blank=True, verbose_name=_("Comment"))
 
+    TypeChoices = TypeChoices
+
     class Meta:
         unique_together = [('org_id', 'name')]
-        verbose_name = _("Command filter rule")
+        verbose_name = _("Command group")
 
     @lazyproperty
     def pattern(self):
         if self.type == 'command':
-            s = self.construct_command_regex(content=self.content)
+            s = self.construct_command_regex(self.content)
         else:
             s = r'{0}'.format(self.content)
         return s
@@ -62,6 +66,17 @@ class CommandGroup(JMSOrgBaseModel):
         s = r'{}'.format('|'.join(regex))
         return s
 
+    def match(self, data):
+        succeed, error, pattern = self.compile_regex(self.pattern, self.ignore_case)
+        if not succeed:
+            return False, ''
+
+        found = pattern.search(data)
+        if not found:
+            return False, ''
+        else:
+            return True, found.group()
+
     @staticmethod
     def compile_regex(regex, ignore_case):
         args = []
@@ -75,27 +90,23 @@ class CommandGroup(JMSOrgBaseModel):
             return False, error, None
         return True, '', pattern
 
-    def match(self, data):
-        succeed, error, pattern = self.compile_regex(self.pattern, self.ignore_case)
-        if not succeed:
-            return False, ''
-
-        found = pattern.search(data)
-        if not found:
-            return False, ''
-        else:
-            return True, found.group()
-
     def __str__(self):
         return '{} % {}'.format(self.type, self.content)
 
 
-class CommandFilterACL(OrgModelMixin, BaseACL):
-    users = models.JSONField(verbose_name=_('User'))
-    assets = models.JSONField(verbose_name=_('Asset'))
-    accounts = models.JSONField(verbose_name=_('Account'))
-    commands = models.ManyToManyField(CommandGroup, verbose_name=_('Commands'))
-    objects = ACLManager.from_queryset(AssetAccountUserACLQuerySet)()
+class CommandFilterACLQuerySet(UserAssetAccountACLQuerySet):
+    def get_command_groups(self):
+        ids = self.values_list('id', flat=True)
+        queryset = CommandFilterACL.command_groups.through.objects.filter(commandfilteracl_id__in=ids)
+        cmd_group_ids = queryset.values_list('commandgroup_id', flat=True)
+        command_groups = CommandGroup.objects.filter(id__in=cmd_group_ids)
+        return command_groups
+
+
+class CommandFilterACL(UserAssetAccountBaseACL):
+    command_groups = models.ManyToManyField(CommandGroup, verbose_name=_('Commands'))
+
+    objects = ACLManager.from_queryset(CommandFilterACLQuerySet)()
 
     class Meta:
         unique_together = ('name', 'org_id')
@@ -122,44 +133,3 @@ class CommandFilterACL(OrgModelMixin, BaseACL):
         assignees = self.reviewers.all()
         ticket.open_by_system(assignees)
         return ticket
-
-    @classmethod
-    def get_command_groups(cls, user_id=None, user_group_id=None, account=None, asset_id=None, org_id=None):
-        # Todo: Do
-        return CommandGroup.objects.all()
-
-        from assets.models import Account, Asset
-        user_groups = []
-        user = get_object_or_none(User, pk=user_id)
-        if user:
-            user_groups.extend(list(user.groups.all()))
-        user_group = get_object_or_none(UserGroup, pk=user_group_id)
-        if user_group:
-            org_id = user_group.org_id
-            user_groups.append(user_group)
-
-        asset = get_object_or_none(Asset, pk=asset_id)
-        q = Q()
-        if user:
-            q |= Q(users=user)
-        if user_groups:
-            q |= Q(user_groups__in=set(user_groups))
-        if account:
-            org_id = account.org_id
-            q |= Q(accounts__contains=account.username) | \
-                 Q(accounts__contains=Account.AliasAccount.ALL)
-        if asset:
-            org_id = asset.org_id
-            q |= Q(assets=asset)
-        if q:
-            cmd_filters = cls.objects.filter(q).filter(is_active=True)
-            if org_id:
-                cmd_filters = cmd_filters.filter(org_id=org_id)
-            filter_ids = cmd_filters.values_list('id', flat=True)
-            command_group_ids = cls.commands.through.objects\
-                .filter(commandfilteracl_id__in=filter_ids)\
-                .values_list('commandgroup_id', flat=True)
-            cmd_groups = CommandGroup.objects.filter(id__in=command_group_ids)
-        else:
-            cmd_groups = CommandGroup.objects.none()
-        return cmd_groups

apps/acls/models/login_acl.py

@@ -9,12 +9,10 @@ from .base import BaseACL
 
 
 class LoginACL(BaseACL):
-    # 用户
     user = models.ForeignKey(
-        'users.User', on_delete=models.CASCADE, verbose_name=_('User'),
-        related_name='login_acls'
+        'users.User', on_delete=models.CASCADE, related_name='login_acls', verbose_name=_('User')
     )
-    # 规则
+    # 规则, ip_group, time_period
     rules = models.JSONField(default=dict, verbose_name=_('Rule'))
 
     class Meta:

apps/acls/models/login_asset_acl.py

@@ -1,17 +1,10 @@
-from django.db import models
 from django.utils.translation import ugettext_lazy as _
 
-from orgs.mixins.models import OrgModelMixin
-from .base import BaseACL, ACLManager, AssetAccountUserACLQuerySet
 
+from .base import UserAssetAccountBaseACL
 
-class LoginAssetACL(BaseACL, OrgModelMixin):
-    # 条件
-    users = models.JSONField(verbose_name=_('User'))
-    accounts = models.JSONField(verbose_name=_('Account'))
-    assets = models.JSONField(verbose_name=_('Asset'))
 
-    objects = ACLManager.from_queryset(AssetAccountUserACLQuerySet)()
+class LoginAssetACL(UserAssetAccountBaseACL):
 
     class Meta:
         unique_together = ('name', 'org_id')

apps/authentication/models/connection_token.py

@@ -135,7 +135,6 @@ class ConnectionToken(OrgModelMixin, JMSBaseModel):
                 'su_from': None,
                 'org_id': self.asset.org_id
             }
-            Account(**data)
         else:
             data = {
                 'name': account.name,
@@ -164,13 +163,12 @@ class ConnectionToken(OrgModelMixin, JMSBaseModel):
     def acl_command_groups(self):
         from acls.models import CommandFilterACL
         kwargs = {
-            'user_id': self.user.id,
+            'user': self.user,
+            'asset': self.asset,
             'account': self.account,
         }
-        if self.asset:
-            kwargs['asset_id'] = self.asset.id
-        cmd_groups = CommandFilterACL.get_command_groups(**kwargs)
-        return cmd_groups
+        command_groups = CommandFilterACL.filter_queryset(**kwargs).get_command_groups()
+        return command_groups
 
 
 class SuperConnectionToken(ConnectionToken):