jumpserver/apps/assets/utils.py

# ~*~ coding: utf-8 ~*~
#
from functools import reduce
from treelib import Tree
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


logger = get_logger(__file__)


def get_system_user_by_name(name):
    system_user = get_object_or_none(SystemUser, name=name)
    return system_user


def get_system_user_by_id(id):
    system_user = get_object_or_none(SystemUser, id=id)
    return system_user


class LabelFilterMixin:
    def get_filter_labels_ids(self):
        query_params = self.request.query_params
        query_keys = query_params.keys()
        all_label_keys = Label.objects.values_list('name', flat=True)
        valid_keys = set(all_label_keys) & set(query_keys)

        if not valid_keys:
            return []

        labels_query = [
            {"name": key, "value": query_params[key]}
            for key in valid_keys
        ]
        args = [Q(**kwargs) for kwargs in labels_query]
        args = reduce(lambda x, y: x | y, args)
        labels_id = Label.objects.filter(args).values_list('id', flat=True)
        return labels_id


class LabelFilter(LabelFilterMixin):
    def filter_queryset(self, queryset):
        queryset = super().filter_queryset(queryset)
        labels_ids = self.get_filter_labels_ids()
        if not labels_ids:
            return queryset
        for labels_id in labels_ids:
            queryset = queryset.filter(labels=labels_id)
        return queryset


class TreeService(Tree):
    tag_sep = ' / '
    cache_key = '_NODE_FULL_TREE'
    cache_time = 3600

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.nodes_assets_map = defaultdict(set)
        self.all_nodes_assets_map = {}
        self.mutex = threading.Lock()

    @classmethod
    @timeit
    def new(cls):
        from .models import Node
        from orgs.utils import get_current_org, set_to_root_org

        origin_org = get_current_org()
        set_to_root_org()
        all_nodes = Node.objects.all()
        origin_org.change_to()

        tree = cls()
        tree.create_node(tag='', identifier='')
        for node in all_nodes:
            tree.create_node(
                tag=node.value, identifier=node.key,
                parent=node.parent_key,
            )
        tree.init_assets_async()
        return tree

    def init_assets_async(self):
        t = threading.Thread(target=self.init_assets)
        t.start()

    def init_assets(self):
        from orgs.utils import get_current_org, set_to_root_org
        with self.mutex:
            origin_org = get_current_org()
            set_to_root_org()
            queryset = Asset.objects.all().valid().values_list('id', 'nodes__key')

            if origin_org:
                origin_org.change_to()
            for asset_id, key in queryset:
                if not key:
                    continue
                self.nodes_assets_map[key].add(asset_id)

    def all_children(self, nid, with_self=True, deep=False):
        children_ids = self.expand_tree(nid)
        if not with_self:
            next(children_ids)
        return [self.get_node(i, deep=deep) for i in children_ids]

    def ancestors(self, nid, with_self=False, deep=False):
        ancestor_ids = list(self.rsearch(nid))
        ancestor_ids.pop()
        if not with_self:
            ancestor_ids.pop(0)
        return [self.get_node(i, deep=deep) for i in ancestor_ids]

    def get_node_full_tag(self, nid):
        ancestors = self.ancestors(nid)
        ancestors.reverse()
        return self.tag_sep.join(n.tag for n in ancestors)

    def get_family(self, nid, deep=False):
        ancestors = self.ancestors(nid, with_self=False, deep=deep)
        children = self.all_children(nid, with_self=False)
        return ancestors + [self[nid]] + children

    def root_node(self):
        return self.get_node(self.root)

    def get_node(self, nid, deep=False):
        node = super().get_node(nid)
        if deep:
            node = self.copy_node(node)
        return node

    def parent(self, nid, deep=False):
        parent = super().parent(nid)
        if deep:
            parent = self.copy_node(parent)
        return parent

    def assets(self, nid):
        with self.mutex:
            assets = self.nodes_assets_map[nid]
            return assets

    def set_assets(self, nid, assets):
        with self.mutex:
            self.nodes_assets_map[nid] = assets

    def all_assets(self, nid):
        assets = self.all_nodes_assets_map.get(nid)
        if assets:
            return assets
        assets = set(self.assets(nid))
        children = self.children(nid)
        for child in children:
            assets.update(self.all_assets(child.identifier))
        return assets

    def assets_amount(self, nid):
        return len(self.all_assets(nid))

    @staticmethod
    def copy_node(node):
        new_node = deepcopy(node)
        new_node.fpointer = None
        return new_node

    def safe_add_ancestors(self, ancestors):
        # 如果祖先节点为1个，那么添加该节点, 父节点是root node
        if len(ancestors) == 1:
            node = ancestors[0]
            parent = self.root_node()
        else:
            node, ancestors = ancestors[0], ancestors[1:]
            parent_id = ancestors[0].identifier
            # 如果父节点不存在, 则先添加父节点
            if not self.contains(parent_id):
                self.safe_add_ancestors(ancestors)
            parent = self.get_node(parent_id)

        print("Add node: {} {}".format(node.identifier, parent.identifier))
        # 如果当前节点已再树中，则移动当前节点到父节点中
        # 这个是由于 当前节点放到了二级节点中
        if self.contains(node.identifier):
            self.move_node(node.identifier, parent.identifier)
        else:
            self.add_node(node, parent)