[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: