perf: 优化权限位 (#8110)

* perf: 优化权限位

* perf: 优化返回的组织

* perf: 保证结果是 ok

* perf: 去掉 distinct

* perf: tree count

Co-authored-by: ibuler <ibuler@qq.com>

apps/rbac/builtin.py

@@ -2,6 +2,10 @@ from django.utils.translation import ugettext_noop
 
 from .const import Scope, system_exclude_permissions, org_exclude_permissions
 
+_view_root_perms = (
+    ('orgs', 'organization', 'view', 'rootorg'),
+)
+
 # 工作台也区分组织后再考虑
 user_perms = (
     ('rbac', 'menupermission', 'view', 'workbench'),
@@ -21,19 +25,23 @@ system_user_perms = (
     ('authentication', 'temptoken', 'add,change,view', 'temptoken'),
     ('authentication', 'accesskey', '*', '*'),
     ('tickets', 'ticket', 'view', 'ticket'),
-    ('orgs', 'organization', 'view', 'rootorg'),
 ) + user_perms
 
-auditor_perms = user_perms + (
+_auditor_perms = (
     ('rbac', 'menupermission', 'view', 'audit'),
     ('audits', '*', '*', '*'),
     ('terminal', 'commandstorage', 'view', 'commandstorage'),
     ('terminal', 'sessionreplay', 'view,download', 'sessionreplay'),
     ('terminal', 'session', '*', '*'),
     ('terminal', 'command', '*', '*'),
-    ('ops', 'commandexecution', 'view', 'commandexecution')
+    ('ops', 'commandexecution', 'view', 'commandexecution'),
 )
 
+auditor_perms = user_perms + _auditor_perms
+
+system_auditor_perms = system_user_perms + _auditor_perms + _view_root_perms
+
+
 app_exclude_perms = [
     ('users', 'user', 'add,delete', 'user'),
     ('orgs', 'org', 'add,delete,change', 'org'),
@@ -101,7 +109,7 @@ class BuiltinRole:
         '1', ugettext_noop('SystemAdmin'), Scope.system, []
     )
     system_auditor = PredefineRole(
-        '2', ugettext_noop('SystemAuditor'), Scope.system, auditor_perms
+        '2', ugettext_noop('SystemAuditor'), Scope.system, system_auditor_perms
     )
     system_component = PredefineRole(
         '4', ugettext_noop('SystemComponent'), Scope.system, app_exclude_perms, 'exclude'

apps/rbac/const.py

@@ -108,8 +108,11 @@ only_system_permissions = (
     ('terminal', 'replaystorage', '*', '*'),
     ('terminal', 'status', '*', '*'),
     ('terminal', 'task', '*', '*'),
+    ('terminal', 'endpoint', '*', '*'),
+    ('terminal', 'endpointrule', '*', '*'),
     ('authentication', '*', '*', '*'),
     ('tickets', '*', '*', '*'),
+    ('orgs', 'organization', 'view', 'rootorg'),
 )
 
 only_org_permissions = (

apps/rbac/models/rolebinding.py

@@ -107,19 +107,23 @@ class RoleBinding(JMSModel):
         roles = Role.get_roles_by_perm(perm)
         with tmp_to_root_org():
             bindings = list(cls.objects.root_all().filter(role__in=roles, user=user))
-        system_bindings = [b for b in bindings if b.scope == Role.Scope.system.value]
 
+        system_bindings = [b for b in bindings if b.scope == Role.Scope.system.value]
+        # 工作台仅限于自己加入的组织
         if perm == 'rbac.view_workbench':
             all_orgs = user.orgs.all()
         else:
             all_orgs = Organization.objects.all()
 
+        # 有系统级别的绑定，就代表在所有组织有这个权限
         if system_bindings:
             orgs = all_orgs
         else:
             org_ids = [b.org.id for b in bindings if b.org]
             orgs = all_orgs.filter(id__in=org_ids)
-        if orgs and user.has_perm('orgs.view_rootorg'):
+
+        # 全局组织
+        if orgs and perm != 'rbac.view_workbench' and user.has_perm('orgs.view_rootorg'):
             orgs = [Organization.root(), *list(orgs)]
         return orgs
 

apps/rbac/tree.py

@@ -1,7 +1,8 @@
 #!/usr/bin/python
 import os
-from collections import defaultdict
 from typing import Callable
+from treelib import Tree
+from treelib.exceptions import NodeIDAbsentError
 
 from django.utils.translation import gettext_lazy as _, gettext, get_language
 from django.conf import settings
@@ -159,6 +160,65 @@ def sort_nodes(node):
     return value
 
 
+class CounterTree(Tree):
+    def get_total_count(self, node):
+        count = getattr(node, '_total_count', None)
+        if count is not None:
+            return count
+
+        if not node.data.isParent:
+            return 1
+
+        count = 0
+        children = self.children(node.identifier)
+        for child in children:
+            if child.data.isParent:
+                count += self.get_total_count(child)
+            else:
+                count += 1
+        node._total_count = count
+        return count
+
+    def get_checked_count(self, node):
+        count = getattr(node, '_checked_count', None)
+        if count is not None:
+            return count
+
+        if not node.data.isParent:
+            if node.data.checked:
+                return 1
+            else:
+                return 0
+
+        count = 0
+        children = self.children(node.identifier)
+        for child in children:
+            if child.data.isParent:
+                count += self.get_checked_count(child)
+            else:
+                if child.data.checked:
+                    count += 1
+        node._checked_count = count
+        return count
+
+    def add_nodes_to_tree(self, ztree_nodes, retry=0):
+        failed = []
+        for node in ztree_nodes:
+            pid = node.pId
+            if retry == 2:
+                pid = '$ROOT$'
+
+            try:
+                self.create_node(node.name, node.id, pid, data=node)
+            except NodeIDAbsentError:
+                failed.append(node)
+        if retry > 2:
+            return
+        if failed:
+            retry += 1
+            return self.add_nodes_to_tree(failed, retry)
+
+
 class PermissionTreeUtil:
     get_permissions: Callable
     action_mapper = {
@@ -183,8 +243,6 @@ class PermissionTreeUtil:
             Permission.get_permissions(scope)
         )
         self.check_disabled = check_disabled
-        self.total_counts = defaultdict(int)
-        self.checked_counts = defaultdict(int)
         self.lang = get_language()
 
     @staticmethod
@@ -211,38 +269,10 @@ class PermissionTreeUtil:
                 'name': name,
                 'pId': view,
             }
-            total_count = self.total_counts[app]
+            node = self._create_node(app_data, 'app', is_open=False)
-            checked_count = self.checked_counts[app]
-            if total_count == 0:
-                continue
-            self.total_counts[view] += total_count
-            self.checked_counts[view] += checked_count
-            node = self._create_node(
-                app_data, total_count, checked_count,
-                'app', is_open=False
-            )
             nodes.append(node)
         return nodes
 
-    def _get_model_counts_mapper(self):
-        model_counts = self.all_permissions \
-            .values('model', 'app', 'content_type') \
-            .order_by('content_type') \
-            .annotate(count=Count('content_type'))
-        model_check_counts = self.permissions \
-            .values('content_type', 'model') \
-            .order_by('content_type') \
-            .annotate(count=Count('content_type'))
-        model_counts_mapper = {
-            i['content_type']: i['count']
-            for i in model_counts
-        }
-        model_check_counts_mapper = {
-            i['content_type']: i['count']
-            for i in model_check_counts
-        }
-        return model_counts_mapper, model_check_counts_mapper
-
     @staticmethod
     def _check_model_xpack(model_id):
         app, model = model_id.split('.', 2)
@@ -263,17 +293,10 @@ class PermissionTreeUtil:
             if not self._check_model_xpack(model_id):
                 continue
 
-            total_count = self.total_counts[model_id]
-            checked_count = self.checked_counts[model_id]
-            if total_count == 0:
-                continue
-
             # 获取 pid
             app = ct.app_label
             if model_id in special_pid_mapper:
                 app = special_pid_mapper[model_id]
-            self.total_counts[app] += total_count
-            self.checked_counts[app] += checked_count
 
             # 获取 name
             name = f'{ct.name}'
@@ -284,7 +307,7 @@ class PermissionTreeUtil:
                 'id': model_id,
                 'name': name,
                 'pId': app,
-            }, total_count, checked_count, 'model', is_open=False)
+            }, 'model', is_open=False)
             nodes.append(node)
         return nodes
 
@@ -334,10 +357,7 @@ class PermissionTreeUtil:
             if title in special_pid_mapper:
                 pid = special_pid_mapper[title]
 
-            self.total_counts[pid] += 1
             checked = p.id in permissions_id
-            if checked:
-                self.checked_counts[pid] += 1
 
             node = TreeNode(**{
                 'id': p.id,
@@ -347,7 +367,7 @@ class PermissionTreeUtil:
                 'isParent': False,
                 'chkDisabled': self.check_disabled,
                 'iconSkin': icon,
-                'checked': p.id in permissions_id,
+                'checked': checked,
                 'open': False,
                 'meta': {
                     'type': 'perm',
@@ -356,13 +376,11 @@ class PermissionTreeUtil:
             nodes.append(node)
         return nodes
 
-    def _create_node(self, data, total_count, checked_count, tp,
+    def _create_node(self, data, tp, is_parent=True, is_open=True, icon='', checked=None):
-                     is_parent=True, is_open=True, icon='', checked=None):
         assert data.get('id')
         assert data.get('name')
         assert data.get('pId') is not None
-        if checked is None:
+
-            checked = total_count == checked_count
         node_data = {
             'isParent': is_parent,
             'iconSkin': icon,
@@ -380,46 +398,58 @@ class PermissionTreeUtil:
             node.name += ('[' + node.id + ']')
         if DEBUG_DB:
             node.name += ('-' + node.id)
-        node.name += f'({checked_count}/{total_count})'
         return node
 
     def _create_root_tree_node(self):
-        total_count = self.all_permissions.count()
+        node = self._create_node(root_node_data, 'root')
-        checked_count = self.permissions.count()
-        node = self._create_node(root_node_data, total_count, checked_count, 'root')
         return node
 
     def _create_views_node(self):
         nodes = []
         for view_data in view_nodes_data:
-            view = view_data['id']
             data = {
                 **view_data,
                 'pId': '$ROOT$',
             }
-            total_count = self.total_counts[view]
+            node = self._create_node(data, 'view', is_open=True)
-            checked_count = self.checked_counts[view]
-            if total_count == 0:
-                continue
-            node = self._create_node(data, total_count, checked_count, 'view', is_open=True)
             nodes.append(node)
         return nodes
 
     def _create_extra_nodes(self):
         nodes = []
         for data in extra_nodes_data:
-            i = data['id']
+            node = self._create_node(data, 'extra', is_open=False)
-            pid = data['pId']
+            nodes.append(node)
-            checked_count = self.checked_counts[i]
+        return nodes
-            total_count = self.total_counts[i]
+
+    @staticmethod
+    def compute_nodes_count(ztree_nodes):
+        tree = CounterTree()
+        reverse_nodes = ztree_nodes[::-1]
+        root = reverse_nodes[0]
+        tree.create_node(root.name, root.id, data=root)
+        tree.add_nodes_to_tree(reverse_nodes[1:])
+        counter_nodes = tree.all_nodes()
+
+        node_counts = {}
+        for n in counter_nodes:
+            if not n:
+                continue
+            total_count = tree.get_total_count(n)
+            checked_count = tree.get_checked_count(n)
+            node_counts[n.identifier] = [checked_count, total_count]
+
+        nodes = []
+        for node in ztree_nodes:
+            counter = node_counts[node.id]
+            if not counter:
+                counter = [0, 0]
+            checked_count, total_count = counter
             if total_count == 0:
                 continue
-            self.total_counts[pid] += total_count
+            node.name += '({}/{})'.format(checked_count, total_count)
-            self.checked_counts[pid] += checked_count
+            if checked_count != 0:
-            node = self._create_node(
+                node.checked = True
-                data, total_count, checked_count,
-                'extra', is_open=False
-            )
             nodes.append(node)
         return nodes
 
@@ -431,5 +461,6 @@ class PermissionTreeUtil:
         nodes += self._create_views_node()
         nodes += [self._create_root_tree_node()]
 
+        nodes = self.compute_nodes_count(nodes)
         nodes.sort(key=sort_nodes)
         return nodes