jumpserver/apps/rbac/ztree/tree.py

import random
from collections import defaultdict
from django.utils.translation import ugettext
from common.tree import TreeNode as RawTreeNode
from django.utils.translation import gettext_lazy as _, gettext
from rbac.models import Permission, ContentType
from django.db.models import F, Count
from .permissions import permission_paths, flag_license_required, flag_sep, flag_scope_system
from .tree_nodes import permission_tree_nodes
from ..const import Scope
from jumpserver.utils import has_valid_xpack_license
from django.conf import settings


class TreeNode(RawTreeNode):
    total_count = 0
    checked_count = 0
    app_label_codename = ''

    def mark_checked_if_need(self):
        if self.isParent:
            self.checked = self.total_count == self.checked_count

    def refresh_name_if_need(self):
        if self.isParent:
            self.name = str(self.name) + f'({self.checked_count}/{self.total_count})'
        elif settings.DEBUG:
            self.name = str(self.name) + f'({self.app_label_codename})'


class TreeNodes:

    def __init__(self):
        self.tree_nodes = defaultdict(TreeNode)

    def add_node(self, data):
        tree_node = self.add(data)
        tree_node.total_count += 1

    def add_leaf(self, data):
        tree_node = self.add(data)
        if not data['checked']:
            return

        parent_node = self.tree_nodes.get(tree_node.pId)
        while parent_node:
            parent_node.checked_count += 1
            parent_node = self.tree_nodes.get(parent_node.pId)

    def add(self, data):
        _id = data['id']
        data['name'] = data.get('name') or data['id']
        tree_node = self.tree_nodes.get(_id, TreeNode(**data))
        self.tree_nodes[tree_node.id] = tree_node
        return tree_node

    def get(self):
        tree_nodes = list(self.tree_nodes.values())
        for tree_node in tree_nodes:
            tree_node.mark_checked_if_need()
            tree_node.refresh_name_if_need()
        return tree_nodes


class ZTree(object):

    has_valid_license = has_valid_xpack_license()

    def __init__(self, checked_permission, scope, check_disabled=False):
        self.scope = scope
        self.checked_permission = self.prefetch_permissions(
            checked_permission
        )
        self.checked_permissions_mapper = {p.id: p for p in self.checked_permission}
        self.permissions = self.prefetch_permissions(
            Permission.get_permissions(scope)
        )
        self.permissions_mapper = {p.app_label_codename: p for p in self.permissions}
        self.content_types_name_mapper = {ct.model: ct.name for ct in ContentType.objects.all()}
        self.check_disabled = check_disabled
        self.tree_nodes = TreeNodes()
        self.show_node_level = 3

    @staticmethod
    def prefetch_permissions(permissions):
        return permissions.select_related('content_type') \
            .annotate(app=F('content_type__app_label')) \
            .annotate(model=F('content_type__model'))

    def get_tree_nodes(self):
        perm_paths = self.__class__.get_permission_paths(self.scope)
        for perm_path in perm_paths:
            self.generate_tree_nodes_by_path(perm_path)
        return self.tree_nodes.get()

    def generate_tree_nodes_by_path(self, perm_path):
        path, perm_app_label_codename = perm_path.rsplit('/', 1)

        # add path
        path_list = path.lstrip('/').split('/')
        pid = ''
        for level, tree_node_id in enumerate(path_list, start=1):
            name = _('Detail') if 'detail' in tree_node_id else tree_node_id
            data = dict({
                'id': tree_node_id,
                'name': name,
                'title': name,
                'pId': pid,
                'isParent': True,
                'chkDisabled': self.check_disabled,
                'open': level < self.show_node_level,
                'meta': {
                    'type': 'perm',
                }
            })
            _data = permission_tree_nodes.get(tree_node_id, {})
            data.update(_data)
            pid = data['id']
            self.tree_nodes.add_node(data)

        # add perm
        if not perm_app_label_codename:
            return
        perm = self.permissions_mapper.get(perm_app_label_codename)
        if perm:
            # 解决同一个权限不能在多个节点的问题
            _id = f'{pid}#{perm.id}'
            name = self._get_permission_name(perm)
            checked = perm.id in self.checked_permissions_mapper
        else:
            #  最终不应该走这里，所有权限都要在数据库里
            _id = perm_app_label_codename
            name = perm_app_label_codename
            checked = False

        data = {
            'id': _id,
            'pId': pid,
            'name': name,
            'title': perm_app_label_codename,
            'chkDisabled': self.check_disabled,
            'app_label_codename': perm_app_label_codename,
            'isParent': False,
            'iconSkin': 'file',
            'open': False,
            'checked': checked,
            'meta': {
                'type': 'perm',
            }
        }
        _data = permission_tree_nodes.get(perm_app_label_codename, {})
        data.update(_data)
        self.tree_nodes.add_leaf(data)

    def _get_permission_name(self, p):
        code_name = p.codename
        action_mapper = {
            'add': ugettext('Create'),
            'view': ugettext('View'),
            'change': ugettext('Update'),
            'delete': ugettext('Delete')
        }
        name = ''
        ct = ''
        if 'add_' in p.codename:
            name = action_mapper['add']
            ct = code_name.replace('add_', '')
        elif 'view_' in p.codename:
            name = action_mapper['view']
            ct = code_name.replace('view_', '')
        elif 'change_' in p.codename:
            name = action_mapper['change']
            ct = code_name.replace('change_', '')
        elif 'delete' in code_name:
            name = action_mapper['delete']
            ct = code_name.replace('delete_', '')

        if ct in self.content_types_name_mapper:
            name += self.content_types_name_mapper[ct]
        else:
            name = gettext(p.name)
            name = name.replace('Can ', '').replace('可以', '')
        return name

    @classmethod
    def get_permissions_app_label_codename(cls, scope):
        perm_paths = cls.get_permission_paths(scope)
        perms = []
        for path in perm_paths:
            path, app_label_code_name = path.rsplit('/', 1)
            if not app_label_code_name:
                continue
            perms.append(app_label_code_name)
        return perms

    @classmethod
    def get_permission_paths(cls, scope):
        perm_paths = []
        for path in permission_paths:
            if flag_sep in path:
                path, flags = path.split(flag_sep)
                if flag_scope_system in flags and scope == Scope.org:
                    continue
                if flag_license_required in flags and not cls.has_valid_license:
                    continue
            perm_paths.append(path)
        return perm_paths