jumpserver/apps/assets/utils.py

# ~*~ coding: utf-8 ~*~
#
from treelib import Tree
from treelib.exceptions import NodeIDAbsentError
from collections import defaultdict
from copy import deepcopy

from common.utils import get_logger, timeit, lazyproperty
from .models import Asset, Node


logger = get_logger(__file__)


class TreeService(Tree):
    tag_sep = ' / '

    @staticmethod
    @timeit
    def get_nodes_assets_map():
        nodes_assets_map = defaultdict(set)
        asset_node_list = Node.assets.through.objects.values_list(
            'asset', 'node__key'
        )
        for asset_id, key in asset_node_list:
            nodes_assets_map[key].add(asset_id)
        return nodes_assets_map

    @classmethod
    @timeit
    def new(cls):
        from .models import Node
        all_nodes = list(Node.objects.all().values("key", "value"))
        all_nodes.sort(key=lambda x: len(x["key"].split(":")))
        tree = cls()
        tree.create_node(tag='', identifier='', data={})
        for node in all_nodes:
            key = node["key"]
            value = node["value"]
            parent_key = ":".join(key.split(":")[:-1])
            tree.safe_create_node(
                tag=value, identifier=key,
                parent=parent_key,
            )
        tree.init_assets()
        return tree

    def init_assets(self):
        node_assets_map = self.get_nodes_assets_map()
        for node in self.all_nodes_itr():
            key = node.identifier
            assets = node_assets_map.get(key, set())
            data = {"assets": assets, "all_assets": None}
            node.data = data

    def safe_create_node(self, **kwargs):
        parent = kwargs.get("parent")
        if not self.contains(parent):
            kwargs['parent'] = self.root
        self.create_node(**kwargs)

    def all_children_ids(self, nid, with_self=True):
        children_ids = self.expand_tree(nid)
        if not with_self:
            next(children_ids)
        return list(children_ids)

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

    def ancestors_ids(self, nid, with_self=True):
        ancestor_ids = list(self.rsearch(nid))
        ancestor_ids.pop()
        if not with_self:
            ancestor_ids.pop(0)
        return ancestor_ids

    def ancestors(self, nid, with_self=False, deep=False, with_assets=True):
        ancestor_ids = self.ancestors_ids(nid, with_self=with_self)
        ancestors = [self.get_node(i, deep=deep) for i in ancestor_ids]
        if with_assets:
            return ancestors
        for n in ancestors:
            n.data['assets'] = set()
            n.data['all_assets'] = None
        return ancestors

    def get_node_full_tag(self, nid):
        ancestors = self.ancestors(nid, with_self=True)
        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

    @staticmethod
    def is_parent(child, parent):
        parent_id = child.bpointer
        return parent_id == parent.identifier

    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

    @lazyproperty
    def invalid_assets(self):
        assets = Asset.objects.filter(is_active=False).values_list('id', flat=True)
        return assets

    def set_assets(self, nid, assets):
        node = self.get_node(nid)
        if node.data is None:
            node.data = {}
        node.data["assets"] = assets

    def assets(self, nid):
        node = self.get_node(nid)
        return node.data.get("assets", set())

    def valid_assets(self, nid):
        return set(self.assets(nid)) - set(self.invalid_assets)

    def all_assets(self, nid):
        node = self.get_node(nid)
        if node.data is None:
            node.data = {}
        all_assets = node.data.get("all_assets")
        if all_assets is not None:
            return all_assets
        all_assets = set(self.assets(nid))
        try:
            children = self.children(nid)
        except NodeIDAbsentError:
            children = []
        for child in children:
            all_assets.update(self.all_assets(child.identifier))
        node.data["all_assets"] = all_assets
        return all_assets

    def all_valid_assets(self, nid):
        return set(self.all_assets(nid)) - set(self.invalid_assets)

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

    def valid_assets_amount(self, nid):
        return len(self.all_valid_assets(nid))

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

    def safe_add_ancestors(self, node, ancestors):
        # 如果没有祖先节点，那么添加该节点, 父节点是root node
        if len(ancestors) == 0:
            parent = self.root_node()
        else:
            parent = ancestors[0]

        # 如果当前节点已再树中，则移动当前节点到父节点中
        # 这个是由于 当前节点放到了二级节点中
        if not self.contains(parent.identifier):
            # logger.debug('Add parent: {}'.format(parent.identifier))
            self.safe_add_ancestors(parent, ancestors[1:])

        if self.contains(node.identifier):
            # msg = 'Move node to parent: {} => {}'.format(
            #     node.identifier, parent.identifier
            # )
            # logger.debug(msg)
            self.move_node(node.identifier, parent.identifier)
        else:
            # logger.debug('Add node: {}'.format(node.identifier))
            self.add_node(node, parent)
    #
    # def __getstate__(self):
    #     self.mutex = None
    #     self.all_nodes_assets_map = {}
    #     self.nodes_assets_map = {}
    #     return self.__dict__

    # def __setstate__(self, state):
    #     self.__dict__ = state