[Update] 修改一些方法

apps/assets/models/asset.py

@@ -152,29 +152,13 @@ class Asset(models.Model):
 
     def get_all_nodes(self, flat=False):
         nodes = []
-        for node in self.get_nodes_or_cache():
+        for node in self.get_nodes():
             _nodes = node.get_ancestor(with_self=True)
             _nodes.append(_nodes)
         if flat:
             nodes = list(reduce(lambda x, y: set(x) | set(y), nodes))
         return nodes
 
-    @property
-    def nodes_cache_key(self):
-        key = "NODES_OF_{}".format(str(self.id))
-        return key
-
-    def get_nodes_or_cache(self):
-        cached = cache.get(self.nodes_cache_key)
-        if cached is not None:
-            return cached
-        nodes = list(self.get_nodes())
-        cache.set(self.nodes_cache_key, nodes, 3600)
-        return nodes
-
-    def expire_nodes_cache(self):
-        cache.delete(self.nodes_cache_key)
-
     @property
     def hardware_info(self):
         if self.cpu_count:

apps/assets/models/node.py

@@ -22,6 +22,21 @@ class Node(models.Model):
     def __str__(self):
         return self.full_value
 
+    def __eq__(self, other):
+        return self.key == other.key
+
+    def __gt__(self, other):
+        if self.is_root():
+            return True
+        self_key = [int(k) for k in self.key.split(':')]
+        other_key = [int(k) for k in other.key.split(':')]
+        if len(self_key) < len(other_key):
+            return True
+        elif len(self_key) > len(other_key):
+            return False
+        else:
+            return self_key[-1] < other_key[-1]
+
     @property
     def name(self):
         return self.value
@@ -93,7 +108,7 @@ class Node(models.Model):
                 Q(nodes__id=self.id) | Q(nodes__isnull=True)
             )
         else:
-            assets = Asset.objects.filter(nodes__id=self.id)
+            assets = self.assets.all()
         return assets
 
     def get_valid_assets(self):
@@ -104,8 +119,8 @@ class Node(models.Model):
         if self.is_root():
             assets = Asset.objects.all()
         else:
-            nodes = self.get_all_children(with_self=True)
-            assets = Asset.objects.filter(nodes__in=nodes).distinct()
+            pattern = r'^{0}$|^{0}:'.format(self.key)
+            assets = Asset.objects.filter(nodes__key__regex=pattern)
         return assets
 
     def get_all_valid_assets(self):
@@ -154,10 +169,3 @@ class Node(models.Model):
         return obj
 
 
-class Tree:
-    def __init__(self, root):
-        self.root = root
-        self.nodes = []
-
-    def add_node(self, node):
-        pass

apps/assets/serializers/asset.py

@@ -16,8 +16,6 @@ class AssetSerializer(BulkSerializerMixin, serializers.ModelSerializer):
     """
     资产的数据结构
     """
-    nodes = serializers.SerializerMethodField()
-
     class Meta:
         model = Asset
         list_serializer_class = BulkListSerializer
@@ -31,10 +29,6 @@ class AssetSerializer(BulkSerializerMixin, serializers.ModelSerializer):
         ])
         return fields
 
-    @staticmethod
-    def get_nodes(obj):
-        return [n.id for n in obj.get_nodes_or_cache()]
-
 
 class AssetGrantedSerializer(serializers.ModelSerializer):
     """

apps/assets/serializers/node.py

@@ -68,7 +68,10 @@ class NodeSerializer(serializers.ModelSerializer):
 
     @staticmethod
     def get_assets_amount(obj):
-        return obj.get_all_assets().count()
+        if obj.is_node:
+            return obj.get_all_assets().count()
+        else:
+            return 0
 
     def get_fields(self):
         fields = super().get_fields()

apps/assets/signals_handler.py

@@ -64,7 +64,6 @@ def on_system_user_assets_change(sender, instance=None, **kwargs):
 @receiver(m2m_changed, sender=Asset.nodes.through)
 def on_asset_node_changed(sender, instance=None, **kwargs):
     if isinstance(instance, Asset):
-        instance.expire_nodes_cache()
         if kwargs['action'] == 'post_add':
             logger.debug("Asset node change signal received")
             nodes = kwargs['model'].objects.filter(pk__in=kwargs['pk_set'])
@@ -81,10 +80,6 @@ def on_asset_node_changed(sender, instance=None, **kwargs):
 def on_node_assets_changed(sender, instance=None, **kwargs):
     if isinstance(instance, Node):
         assets = kwargs['model'].objects.filter(pk__in=kwargs['pk_set'])
-        # 清理资产节点缓存
-        for asset in assets:
-            asset.expire_nodes_cache()
-
         if kwargs['action'] == 'post_add':
             logger.debug("Node assets change signal received")
             # 重新关联系统用户和资产的关系

apps/perms/api.py

@@ -70,7 +70,8 @@ class UserGrantedAssetsApi(ListAPIView):
         else:
             user = self.request.user
 
-        for k, v in AssetPermissionUtil.get_user_assets(user).items():
+        util = AssetPermissionUtil(user)
+        for k, v in util.get_assets().items():
             if k.is_unixlike():
                 system_users_granted = [s for s in v if s.protocol == 'ssh']
             else:
@@ -95,7 +96,8 @@ class UserGrantedNodesApi(ListAPIView):
             user = get_object_or_404(User, id=user_id)
         else:
             user = self.request.user
-        nodes = AssetPermissionUtil.get_user_nodes_with_assets(user)
+        util = AssetPermissionUtil(user)
+        nodes = util.get_nodes_with_assets()
         return nodes.keys()
 
     def get_permissions(self):
@@ -116,7 +118,8 @@ class UserGrantedNodesWithAssetsApi(ListAPIView):
         else:
             user = get_object_or_404(User, id=user_id)
 
-        nodes = AssetPermissionUtil.get_user_nodes_with_assets(user)
+        util = AssetPermissionUtil(user)
+        nodes = util.get_nodes_with_assets()
         for node, _assets in nodes.items():
             assets = _assets.keys()
             for k, v in _assets.items():
@@ -147,8 +150,9 @@ class UserGrantedNodeAssetsApi(ListAPIView):
             user = get_object_or_404(User, id=user_id)
         else:
             user = self.request.user
+        util = AssetPermissionUtil(user)
         node = get_object_or_404(Node, id=node_id)
-        nodes = AssetPermissionUtil.get_user_nodes_with_assets(user)
+        nodes = util.get_nodes_with_assets()
         assets = nodes.get(node, [])
         for asset, system_users in assets.items():
             asset.system_users_granted = system_users
@@ -172,7 +176,8 @@ class UserGroupGrantedAssetsApi(ListAPIView):
             return queryset
 
         user_group = get_object_or_404(UserGroup, id=user_group_id)
-        assets = AssetPermissionUtil.get_user_group_assets(user_group)
+        util = AssetPermissionUtil(user_group)
+        assets = util.get_assets()
         for k, v in assets.items():
             k.system_users_granted = v
             queryset.append(k)
@@ -189,7 +194,8 @@ class UserGroupGrantedNodesApi(ListAPIView):
 
         if group_id:
             group = get_object_or_404(UserGroup, id=group_id)
-            nodes = AssetPermissionUtil.get_user_group_nodes_with_assets(group)
+            util = AssetPermissionUtil(group)
+            nodes = util.get_nodes_with_assets()
             return nodes.keys()
         return queryset
 
@@ -206,7 +212,8 @@ class UserGroupGrantedNodesWithAssetsApi(ListAPIView):
             return queryset
 
         user_group = get_object_or_404(UserGroup, id=user_group_id)
-        nodes = AssetPermissionUtil.get_user_group_nodes_with_assets(user_group)
+        util = AssetPermissionUtil(user_group)
+        nodes = util.get_nodes_with_assets()
         for node, _assets in nodes.items():
             assets = _assets.keys()
             for asset, system_users in _assets.items():
@@ -226,7 +233,8 @@ class UserGroupGrantedNodeAssetsApi(ListAPIView):
 
         user_group = get_object_or_404(UserGroup, id=user_group_id)
         node = get_object_or_404(Node, id=node_id)
-        nodes = AssetPermissionUtil.get_user_group_nodes_with_assets(user_group)
+        util = AssetPermissionUtil(user_group)
+        nodes = util.get_nodes_with_assets()
         assets = nodes.get(node, [])
         for asset, system_users in assets.items():
             asset.system_users_granted = system_users
@@ -246,7 +254,8 @@ class ValidateUserAssetPermissionView(APIView):
         asset = get_object_or_404(Asset, id=asset_id)
         system_user = get_object_or_404(SystemUser, id=system_id)
 
-        assets_granted = AssetPermissionUtil.get_user_assets(user)
+        util = AssetPermissionUtil(user)
+        assets_granted = util.get_assets()
         if system_user in assets_granted.get(asset, []):
             return Response({'msg': True}, status=200)
         else:

apps/perms/utils.py

@@ -1,301 +1,118 @@
 # coding: utf-8
 
 from __future__ import absolute_import, unicode_literals
-import collections
 from collections import defaultdict
 from django.db.models import Q
-from django.utils import timezone
-import copy
 
-from common.utils import set_or_append_attr_bulk, get_logger
+from common.utils import get_logger
 from .models import AssetPermission
-from .hands import Node, User, UserGroup, Asset, SystemUser
+from .hands import Node
 
 logger = get_logger(__file__)
 
 
-class Tree:
-    def __init__(self):
-        self.__all_nodes = list(Node.objects.all().prefetch_related('assets'))
-        self.__node_asset_map = defaultdict(set)
-        self.nodes = defaultdict(dict)
-        self.root = Node.root()
-        self.init_node_asset_map()
+def get_user_permissions(user, include_group=True):
+    if include_group:
+        groups = user.groups.all()
+        arg = Q(users=user) | Q(user_groups=groups)
+    else:
+        arg = Q(users=user)
+    return AssetPermission.objects.all().valid().filter(arg)
 
-    def init_node_asset_map(self):
-        for node in self.__all_nodes:
-            assets = node.get_assets().values_list('id', flat=True)
-            for asset in assets:
-                self.__node_asset_map[str(asset)].add(node)
 
-    def add_asset(self, asset, system_users):
-        nodes = self.__node_asset_map.get(str(asset.id), [])
-        self.add_nodes(nodes)
-        for node in nodes:
-            self.nodes[node][asset].update(system_users)
+def get_user_group_permissions(user_group):
+    return AssetPermission.objects.all().valid().filter(
+        user_groups=user_group
+    )
 
-    def add_node(self, node):
-        if node in self.nodes:
-            return
-        else:
-            self.nodes[node] = defaultdict(set)
-        if node.key == self.root.key:
-            return
-        parent_key = ':'.join(node.key.split(':')[:-1])
-        for n in self.__all_nodes:
-            if n.key == parent_key:
-                self.add_node(n)
-                break
 
-    def add_nodes(self, nodes):
-        for node in nodes:
-            self.add_node(node)
+def get_asset_permissions(asset, include_node=True):
+    if include_node:
+        nodes = asset.get_all_nodes(flat=True)
+        arg = Q(assets=asset) | Q(nodes=nodes)
+    else:
+        arg = Q(assets=asset)
+    return AssetPermission.objects.all().valid().filter(arg)
+
+
+def get_node_permissions(node):
+    return AssetPermission.objects.all().valid().filter(nodes=node)
+
+
+def get_system_user_permissions(system_user):
+    return AssetPermission.objects.valid().all().filter(
+        system_users=system_user
+    )
 
 
 class AssetPermissionUtil:
-    @staticmethod
-    def get_user_permissions(user):
-        return AssetPermission.objects.all().valid().filter(users=user)
+    get_permissions_map = {
+        "User": get_user_permissions,
+        "UserGroup": get_user_group_permissions,
+        "Asset": get_asset_permissions,
+        "Node": get_node_permissions,
+        "SystemUser": get_node_permissions,
+    }
 
-    @staticmethod
-    def get_user_group_permissions(user_group):
-        return AssetPermission.objects.all().valid().filter(
-            user_groups=user_group
-        )
-
-    @staticmethod
-    def get_asset_permissions(asset):
-        return AssetPermission.objects.all().valid().filter(
-            assets=asset
-        )
-
-    @staticmethod
-    def get_node_permissions(node):
-        return AssetPermission.objects.all().valid().filter(nodes=node)
-
-    @staticmethod
-    def get_system_user_permissions(system_user):
-        return AssetPermission.objects.valid().all().filter(
-            system_users=system_user
-        )
-
-    @classmethod
-    def get_user_group_nodes(cls, group):
-        nodes = defaultdict(set)
-        permissions = cls.get_user_group_permissions(group)
-        for perm in permissions:
-            _nodes = perm.nodes.all()
-            _system_users = perm.system_users.all()
-            set_or_append_attr_bulk(_nodes, 'permission', perm.id)
-            for node in _nodes:
-                nodes[node].update(set(_system_users))
-        return nodes
-
-    @classmethod
-    def get_user_group_assets_direct(cls, group):
-        assets = defaultdict(set)
-        permissions = cls.get_user_group_permissions(group)
-        for perm in permissions:
-            _assets = perm.assets.all().valid()
-            _system_users = perm.system_users.all()
-            set_or_append_attr_bulk(_assets, 'permission', perm.id)
-            for asset in _assets:
-                assets[asset].update(set(_system_users))
-        return assets
-
-    @classmethod
-    def get_user_group_nodes_assets(cls, group):
-        assets = defaultdict(set)
-        nodes = cls.get_user_group_nodes(group)
-        for node, _system_users in nodes.items():
-            _assets = node.get_all_valid_assets()
-            set_or_append_attr_bulk(_assets, 'inherit_node', node.id)
-            set_or_append_attr_bulk(_assets, 'permission', getattr(node, 'permission', None))
-            for asset in _assets:
-                assets[asset].update(set(_system_users))
-        return assets
-
-    @classmethod
-    def get_user_group_assets(cls, group):
-        assets = defaultdict(set)
-        _assets = cls.get_user_group_assets_direct(group)
-        _nodes_assets = cls.get_user_group_nodes_assets(group)
-        for asset, _system_users in _assets.items():
-            assets[asset].update(set(_system_users))
-        for asset, _system_users in _nodes_assets.items():
-            assets[asset].update(set(_system_users))
-        return assets
-
-    @classmethod
-    def get_user_group_nodes_with_assets(cls, group):
-        """
-        :param group:
-        :return: {node: {asset: set(su1, su2)}}
-        """
-        _assets = cls.get_user_group_assets(group)
-        tree = Tree()
-        for asset, _system_users in _assets.items():
-            _nodes = asset.get_nodes_or_cache()
-            tree.add_nodes(_nodes)
-            for node in _nodes:
-                tree.nodes[node][asset].update(_system_users)
-        return tree.nodes
-
-    @classmethod
-    def get_user_assets_direct(cls, user):
-        assets = defaultdict(set)
-        permissions = list(cls.get_user_permissions(user))
-        for perm in permissions:
-            _assets = perm.assets.all().valid()
-            _system_users = perm.system_users.all()
-            set_or_append_attr_bulk(_assets, 'permission', perm.id)
-            for asset in _assets:
-                assets[asset].update(set(_system_users))
-        return assets
-
-    @classmethod
-    def get_user_nodes_direct(cls, user):
-        nodes = defaultdict(set)
-        permissions = cls.get_user_permissions(user)
-        for perm in permissions:
-            _nodes = perm.nodes.all()
-            _system_users = perm.system_users.all()
-            set_or_append_attr_bulk(_nodes, 'permission', perm.id)
-            for node in _nodes:
-                nodes[node].update(set(_system_users))
-        return nodes
-
-    @classmethod
-    def get_user_nodes_inherit_group(cls, user):
-        nodes = defaultdict(set)
-        groups = user.groups.all()
-        for group in groups:
-            _nodes = cls.get_user_group_nodes(group)
-            for node, system_users in _nodes.items():
-                nodes[node].update(set(system_users))
-        return nodes
-
-    @classmethod
-    def get_user_nodes(cls, user):
-        nodes = cls.get_user_nodes_direct(user)
-        nodes_inherit = cls.get_user_nodes_inherit_group(user)
-        for node, system_users in nodes_inherit.items():
-            nodes[node].update(set(system_users))
-        return nodes
-
-    @classmethod
-    def get_user_nodes_assets_direct(cls, user):
-        assets = defaultdict(set)
-        nodes = cls.get_user_nodes_direct(user)
-        for node, _system_users in nodes.items():
-            _assets = node.get_all_valid_assets()
-            set_or_append_attr_bulk(_assets, 'inherit_node', node.id)
-            set_or_append_attr_bulk(_assets, 'permission', getattr(node, 'permission', None))
-            for asset in _assets:
-                assets[asset].update(set(_system_users))
-        return assets
-
-    @classmethod
-    def get_user_assets_inherit_group(cls, user):
-        assets = defaultdict(set)
-        for group in user.groups.all():
-            _assets = cls.get_user_group_assets(group)
-            set_or_append_attr_bulk(_assets, 'inherit_group', group.id)
-            for asset, _system_users in _assets.items():
-                assets[asset].update(_system_users)
-        return assets
-
-    @classmethod
-    def get_user_assets(cls, user):
-        assets = defaultdict(set)
-        _assets_direct = cls.get_user_assets_direct(user)
-        _nodes_assets_direct = cls.get_user_nodes_assets_direct(user)
-        _assets_inherit_group = cls.get_user_assets_inherit_group(user)
-        for asset, _system_users in _assets_direct.items():
-            assets[asset].update(_system_users)
-        for asset, _system_users in _nodes_assets_direct.items():
-            assets[asset].update(_system_users)
-        for asset, _system_users in _assets_inherit_group.items():
-            assets[asset].update(_system_users)
-        return assets
-
-    @classmethod
-    def get_user_nodes_with_assets(cls, user):
-        """
-        :param user:
-        :return: {node: {asset: set(su1, su2)}}
-        """
-        tree = Tree()
-        _assets = cls.get_user_assets(user)
-        for asset, _system_users in _assets.items():
-            tree.add_asset(asset, _system_users)
-            # _nodes = asset.get_nodes()
-            # tree.add_nodes(_nodes)
-            # for node in _nodes:
-            #     tree.nodes[node][asset].update(_system_users)
-        return tree.nodes
-
-    @classmethod
-    def get_system_user_assets(cls, system_user):
-        assets = set()
-        permissions = cls.get_system_user_permissions(system_user)
-        for perm in permissions:
-            assets.update(set(perm.assets.all().valid()))
-            nodes = perm.nodes.all()
-            for node in nodes:
-                assets.update(set(node.get_all_valid_assets()))
-        return assets
-
-    @classmethod
-    def get_node_system_users(cls, node):
-        system_users = set()
-        permissions = cls.get_node_permissions(node)
-        for perm in permissions:
-            system_users.update(perm.system_users.all())
-        return system_users
-
-
-class AssetPermissionUtilsV2:
     def __init__(self, obj):
         self.object = obj
         self._permissions = None
 
-    @staticmethod
-    def get_user_permissions(user):
-        groups = user.groups.all()
-        return AssetPermission.objects.all().valid().filter(
-            Q(users=user) | Q(user_groups=groups)
-        )
-
-    @staticmethod
-    def get_user_group_permissions(user_group):
-        return AssetPermission.objects.all().valid().filter(
-            user_groups=user_group
-        )
-
-    @staticmethod
-    def get_asset_permissions(asset):
-        direct_nodes = asset.get_nodes_or_cache()
-
-        return AssetPermission.objects.all().valid().filter(
-            Q(assets=asset) | Q(nodes=direct_nodes)
-        )
-
-    @staticmethod
-    def get_node_permissions(node):
-        return AssetPermission.objects.all().valid().filter(nodes=node)
-
-    @staticmethod
-    def get_system_user_permissions(system_user):
-        return AssetPermission.objects.valid().all().filter(
-            system_users=system_user
-        )
-
     @property
     def permissions(self):
         if self._permissions:
             return self._permissions
-        if isinstance(self.object, User):
-            pass
+        object_cls = self.object.__class__.__name__
+        func = self.get_permissions_map[object_cls]
+        permissions = func(self.object)
+        self._permissions = permissions
+        return permissions
+
+    def get_nodes_direct(self):
+        """
+        返回用户/组授权规则直接关联的节点
+        :return: {node1: set(system_user1,)}
+        """
+        nodes = defaultdict(set)
+        permissions = self.permissions.prefetch_related('nodes', 'system_users')
+        for perm in permissions:
+            for node in perm.nodes.all():
+                nodes[node].update(perm.system_users.all())
+        return nodes
+
+    def get_assets_direct(self):
+        """
+        返回用户授权规则直接关联的资产
+        :return: {asset1: set(system_user1,)}
+        """
+        assets = defaultdict(set)
+        permissions = self.permissions.prefetch_related('assets', 'system_users')
+        for perm in permissions:
+            for asset in perm.assets.all().valid().prefetch_related('nodes'):
+                assets[asset].update(perm.system_users.all())
+        return assets
+
+    def get_assets(self):
+        assets = self.get_assets_direct()
+        nodes = self.get_nodes_direct()
+        for node, system_users in nodes.items():
+            _assets = node.get_all_assets().valid().prefetch_related('nodes')
+            for asset in _assets:
+                if isinstance(asset, Node):
+                    print(_assets)
+                assets[asset].update(system_users)
+        return assets
+
+    def get_nodes_with_assets(self):
+        assets = self.get_assets()
+        nodes = defaultdict(dict)
+        for asset, system_users in assets.items():
+            _nodes = asset.nodes.all()
+            for node in _nodes:
+                if asset in nodes[node]:
+                    nodes[node][asset].update(system_users)
+                else:
+                    nodes[node] = {asset: system_users}
+        return nodes