From e61bae5ee4140d4fe5bbd4f3ca688bcd0d6e8a87 Mon Sep 17 00:00:00 2001 From: fit2bot <68588906+fit2bot@users.noreply.github.com> Date: Wed, 20 Apr 2022 18:50:53 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=E6=9D=83=E9=99=90?= =?UTF-8?q?=E4=BD=8D=20(#8110)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf: 优化权限位 * perf: 优化返回的组织 * perf: 保证结果是 ok * perf: 去掉 distinct * perf: tree count Co-authored-by: ibuler --- apps/rbac/builtin.py | 16 ++- apps/rbac/const.py | 3 + apps/rbac/models/rolebinding.py | 8 +- apps/rbac/tree.py | 167 +++++++++++++++++++------------- 4 files changed, 120 insertions(+), 74 deletions(-) diff --git a/apps/rbac/builtin.py b/apps/rbac/builtin.py index 179889111..b736ba226 100644 --- a/apps/rbac/builtin.py +++ b/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' diff --git a/apps/rbac/const.py b/apps/rbac/const.py index d9b80b78a..b84b7c69e 100644 --- a/apps/rbac/const.py +++ b/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 = ( diff --git a/apps/rbac/models/rolebinding.py b/apps/rbac/models/rolebinding.py index dc09f75d2..c0ac806ef 100644 --- a/apps/rbac/models/rolebinding.py +++ b/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 diff --git a/apps/rbac/tree.py b/apps/rbac/tree.py index a585bdf5c..bae0b930e 100644 --- a/apps/rbac/tree.py +++ b/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] - 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 - ) + node = self._create_node(app_data, '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, - is_parent=True, is_open=True, icon='', checked=None): + def _create_node(self, data, tp, 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() - checked_count = self.permissions.count() - node = self._create_node(root_node_data, total_count, checked_count, 'root') + node = self._create_node(root_node_data, '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] - checked_count = self.checked_counts[view] - if total_count == 0: - continue - node = self._create_node(data, total_count, checked_count, 'view', is_open=True) + node = self._create_node(data, 'view', is_open=True) nodes.append(node) return nodes def _create_extra_nodes(self): nodes = [] for data in extra_nodes_data: - i = data['id'] - pid = data['pId'] - checked_count = self.checked_counts[i] - total_count = self.total_counts[i] + node = self._create_node(data, 'extra', is_open=False) + nodes.append(node) + return nodes + + @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 - self.checked_counts[pid] += checked_count - node = self._create_node( - data, total_count, checked_count, - 'extra', is_open=False - ) + node.name += '({}/{})'.format(checked_count, total_count) + if checked_count != 0: + node.checked = True 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