From 58e36b5f635e0dad38e1272282387de217a57221 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Wed, 6 Jun 2018 11:13:13 +0800
Subject: [PATCH] =?UTF-8?q?[Bugfix]=20=E6=9A=82=E6=97=B6=E8=BF=98=E4=B8=8D?=
 =?UTF-8?q?=E8=83=BD=E5=88=A0=E9=99=A4Tree=E8=BF=99=E4=B8=AA=E7=B1=BB?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/perms/utils.py | 55 +++++++++++++++++++++++++++++++++++++--------
 1 file changed, 46 insertions(+), 9 deletions(-)

diff --git a/apps/perms/utils.py b/apps/perms/utils.py
index 7ff64b00d..a76953aac 100644
--- a/apps/perms/utils.py
+++ b/apps/perms/utils.py
@@ -11,6 +11,44 @@ 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 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 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_user_permissions(user, include_group=True):
     if include_group:
         groups = user.groups.all()
@@ -57,6 +95,7 @@ class AssetPermissionUtil:
     def __init__(self, obj):
         self.object = obj
         self._permissions = None
+        self._assets = None
 
     @property
     def permissions(self):
@@ -93,6 +132,8 @@ class AssetPermissionUtil:
         return assets
 
     def get_assets(self):
+        if self._assets:
+            return self._assets
         assets = self.get_assets_direct()
         nodes = self.get_nodes_direct()
         for node, system_users in nodes.items():
@@ -101,7 +142,8 @@ class AssetPermissionUtil:
                 if isinstance(asset, Node):
                     print(_assets)
                 assets[asset].update(system_users)
-        return assets
+        self._assets = assets
+        return self._assets
 
     def get_nodes_with_assets(self):
         """
@@ -110,14 +152,9 @@ class AssetPermissionUtil:
         :return:
         """
         assets = self.get_assets()
-        nodes = defaultdict(dict)
+        tree = Tree()
         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
+            tree.add_asset(asset, system_users)
+        return tree.nodes