Merge branch 'dev' of github.com:jumpserver/jumpserver into dev

apps/assets/utils.py

@@ -1,14 +1,12 @@
 # ~*~ coding: utf-8 ~*~
 #
-from functools import reduce
 from treelib import Tree
+from treelib.exceptions import NodeIDAbsentError
 from collections import defaultdict
 from copy import deepcopy
-import threading
-from django.db.models import Q
 
 from common.utils import get_object_or_none, get_logger, timeit
-from .models import SystemUser, Label, Asset
+from .models import SystemUser, Asset
 
 
 logger = get_logger(__file__)
@@ -31,6 +29,7 @@ class TreeService(Tree):
         super().__init__(*args, **kwargs)
         self.nodes_assets_map = defaultdict(set)
         self.all_nodes_assets_map = {}
+        self._invalid_assets = frozenset()
 
     @classmethod
     @timeit
@@ -47,25 +46,34 @@ class TreeService(Tree):
                 key = node["key"]
                 value = node["value"]
                 parent_key = ":".join(key.split(":")[:-1])
-                tree.create_node(
+                tree.safe_create_node(
                     tag=value, identifier=key,
                     parent=parent_key,
                 )
             tree.init_assets()
         return tree
 
+    @timeit
     def init_assets(self):
         from orgs.utils import tmp_to_root_org
         self.all_nodes_assets_map = {}
         self.nodes_assets_map = defaultdict(set)
-        logger.debug('Init tree assets')
         with tmp_to_root_org():
             queryset = Asset.objects.all().values_list('id', 'nodes__key')
+            invalid_assets = Asset.objects.filter(is_active=False)\
+                .values_list('id', flat=True)
+            self._invalid_assets = frozenset(invalid_assets)
             for asset_id, key in queryset:
                 if not key:
                     continue
                 self.nodes_assets_map[key].add(asset_id)
 
+    def safe_create_node(self, **kwargs):
+        parent = kwargs.get("parent")
+        if not self.contains(parent):
+            kwargs['parent'] = self.root
+        self.create_node(**kwargs)
+
     def all_children_ids(self, nid, with_self=True):
         children_ids = self.expand_tree(nid)
         if not with_self:
@@ -93,6 +101,11 @@ class TreeService(Tree):
         children = self.all_children(nid, with_self=False)
         return ancestors + [self[nid]] + children
 
+    @staticmethod
+    def is_parent(child, parent):
+        parent_id = child.bpointer
+        return parent_id == parent.identifier
+
     def root_node(self):
         return self.get_node(self.root)
 
@@ -108,12 +121,15 @@ class TreeService(Tree):
             parent = self.copy_node(parent)
         return parent
 
+    def set_assets(self, nid, assets):
+        self.nodes_assets_map[nid] = set(assets)
+
     def assets(self, nid):
         assets = self.nodes_assets_map[nid]
         return assets
 
-    def set_assets(self, nid, assets):
+    def valid_assets(self, nid):
-        self.nodes_assets_map[nid] = assets
+        return set(self.assets(nid)) - set(self._invalid_assets)
 
     def all_assets(self, nid):
         assets = self.all_nodes_assets_map.get(nid)
@@ -123,41 +139,55 @@ class TreeService(Tree):
         children = self.children(nid)
         for child in children:
             assets.update(self.all_assets(child.identifier))
+        self.all_nodes_assets_map[nid] = assets
         return assets
 
+    def all_valid_assets(self, nid):
+        return set(self.all_assets(nid)) - set(self._invalid_assets)
+
     def assets_amount(self, nid):
         return len(self.all_assets(nid))
 
+    def valid_assets_amount(self, nid):
+        return len(self.all_valid_assets(nid))
+
     @staticmethod
     def copy_node(node):
         new_node = deepcopy(node)
         new_node.fpointer = None
         return new_node
 
-    def safe_add_ancestors(self, ancestors):
+    def safe_add_ancestors(self, node, ancestors):
-        # 如果祖先节点为1个，那么添加该节点, 父节点是root node
+        # 如果没有祖先节点，那么添加该节点, 父节点是root node
-        if len(ancestors) == 1:
+        if len(ancestors) == 0:
-            node = ancestors[0]
             parent = self.root_node()
         else:
-            node, ancestors = ancestors[0], ancestors[1:]
+            parent = ancestors[0]
-            parent_id = ancestors[0].identifier
-            # 如果父节点不存在, 则先添加父节点
-            if not self.contains(parent_id):
-                self.safe_add_ancestors(ancestors)
-            parent = self.get_node(parent_id)
 
         # 如果当前节点已再树中，则移动当前节点到父节点中
         # 这个是由于 当前节点放到了二级节点中
+        if not self.contains(parent.identifier):
+            # logger.debug('Add parent: {}'.format(parent.identifier))
+            self.safe_add_ancestors(parent, ancestors[1:])
+
         if self.contains(node.identifier):
+            # msg = 'Move node to parent: {} => {}'.format(
+            #     node.identifier, parent.identifier
+            # )
+            # logger.debug(msg)
             self.move_node(node.identifier, parent.identifier)
         else:
+            # logger.debug('Add node: {}'.format(node.identifier))
             self.add_node(node, parent)
+
     #
     # def __getstate__(self):
     #     self.mutex = None
     #     return self.__dict__
     #
-    # def __setstate__(self, state):
+
-    #     self.__dict__ = state
+    def __setstate__(self, state):
+        self.__dict__ = state
+        if '_invalid_assets' not in state:
+            self._invalid_assets = frozenset()
         # self.mutex = threading.Lock()

apps/common/utils/common.py

@@ -194,7 +194,7 @@ def timeit(func):
         now = time.time()
         result = func(*args, **kwargs)
         using = (time.time() - now) * 1000
-        msg = "Call {} end, using: {:.1f}ms".format(func.__name__, using)
+        msg = "End call {}, using: {:.1f}ms".format(func.__name__, using)
         logger.debug(msg)
         return result
     return wrapper

apps/perms/api/user_permission/mixin.py

@@ -33,7 +33,7 @@ class UserNodeTreeMixin:
         _queryset = []
 
         for node in nodes:
-            assets_amount = self.tree.assets_amount(node.key)
+            assets_amount = self.tree.valid_assets_amount(node.key)
             if assets_amount == 0 and node.key != Node.empty_key:
                 continue
             node.assets_amount = assets_amount

apps/perms/utils/asset_permission.py

@@ -1,11 +1,8 @@
 # coding: utf-8
 
-import time
 import pickle
 from collections import defaultdict
 from functools import reduce
-from hashlib import md5
-import json
 
 from django.core.cache import cache
 from django.db.models import Q
@@ -94,6 +91,7 @@ class AssetPermissionUtilCacheMixin:
         user_tree = pickle.loads(data)
         return user_tree
 
+    @timeit
     def get_user_tree_from_cache_if_need(self):
         if not self.user_tree_cache_enable:
             return None
@@ -278,9 +276,9 @@ class AssetPermissionUtilV2(AssetPermissionUtilCacheMixin):
             )
             if not ancestors:
                 continue
-            parent_id = ancestors[0].identifier
+            user_tree.safe_add_ancestors(child, ancestors)
-            user_tree.safe_add_ancestors(ancestors)
+            # parent_id = ancestors[0].identifier
-            user_tree.move_node(child.identifier, parent_id)
+            # user_tree.move_node(child.identifier, parent_id)
 
     @staticmethod
     def add_empty_node_if_need(user_tree):
@@ -315,6 +313,7 @@ class AssetPermissionUtilV2(AssetPermissionUtilCacheMixin):
             self.set_user_tree_to_local(user_tree)
             return user_tree
         user_tree = TreeService()
+        user_tree._invalid_assets = self.full_tree._invalid_assets
         full_tree_root = self.full_tree.root_node()
         user_tree.create_node(
             tag=full_tree_root.tag,

apps/users/templates/users/_granted_assets.html

@@ -103,6 +103,7 @@ function onSelected(event, treeNode) {
 
 
 function initTree(refresh) {
+    var asyncUrl = setUrlParam(treeUrl, 'cache_policy', '1');
     var setting = {
         view: {
             dblClickExpand: false,
@@ -115,7 +116,7 @@ function initTree(refresh) {
         },
         async: {
 			enable: true,
-			url: treeUrl,
+			url: asyncUrl,
 			autoParam: ["id=key", "name=n", "level=lv"],
             type: 'get'
 		},