jumpserver/apps/assets/utils.py

# ~*~ coding: utf-8 ~*~
#
import time
from django.db.models import Prefetch

from common.utils import get_object_or_none, get_logger
from common.struct import Stack
from .models import SystemUser, Label, Node, 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 LabelFilter:
    def filter_queryset(self, queryset):
        queryset = super().filter_queryset(queryset)
        query_keys = self.request.query_params.keys()
        all_label_keys = Label.objects.values_list('name', flat=True)
        valid_keys = set(all_label_keys) & set(query_keys)
        labels_query = {}
        for key in valid_keys:
            labels_query[key] = self.request.query_params.get(key)

        conditions = []
        for k, v in labels_query.items():
            query = {'labels__name': k, 'labels__value': v}
            conditions.append(query)

        if conditions:
            for kwargs in conditions:
                queryset = queryset.filter(**kwargs)
        return queryset


class NodeUtil:
    def __init__(self, with_assets_amount=False, debug=False):
        self.stack = Stack()
        self._nodes = {}
        self.with_assets_amount = with_assets_amount
        self._debug = debug
        self.init()

    @staticmethod
    def sorted_by(node):
        return [int(i) for i in node.key.split(':')]

    def get_queryset(self):
        all_nodes = Node.objects.all()
        if self.with_assets_amount:
            all_nodes = all_nodes.prefetch_related(
                Prefetch('assets', queryset=Asset.objects.all().only('id'))
            )
            all_nodes = list(all_nodes)
            for node in all_nodes:
                node._assets = set(node.assets.all())
        return all_nodes

    def get_all_nodes(self):
        all_nodes = sorted(self.get_queryset(), key=self.sorted_by)

        guarder = Node(key='', value='Guarder')
        guarder._assets = []
        all_nodes.append(guarder)
        return all_nodes

    def push_to_stack(self, node):
        # 入栈之前检查
        # 如果栈是空的，证明是一颗树的根部
        if self.stack.is_empty():
            node._full_value = node.value
            node._parents = []
        else:
            # 如果不是根节点,
            # 该节点的祖先应该是父节点的祖先加上父节点
            # 该节点的名字是父节点的名字+自己的名字
            node._parents = [self.stack.top] + self.stack.top._parents
            node._full_value = ' / '.join(
                [self.stack.top._full_value, node.value]
            )
        node._children = []
        node._all_children = []
        self.debug("入栈: {}".format(node.key))
        self.stack.push(node)

    # 出栈
    def pop_from_stack(self):
        _node = self.stack.pop()
        self.debug("出栈: {} 栈顶: {}".format(_node.key, self.stack.top.key if self.stack.top else None))
        self._nodes[_node.key] = _node
        if not self.stack.top:
            return
        if self.with_assets_amount:
            self.stack.top._assets.update(_node._assets)
            _node._assets_amount = len(_node._assets)
            delattr(_node, '_assets')
        self.stack.top._children.append(_node)
        self.stack.top._all_children.extend([_node] + _node._children)

    def init(self):
        all_nodes = self.get_all_nodes()
        for node in all_nodes:
            self.debug("准备: {} 栈顶: {}".format(node.key, self.stack.top.key if self.stack.top else None))
            # 入栈之前检查，该节点是不是栈顶节点的子节点
            # 如果不是，则栈顶出栈
            while self.stack.top and not self.stack.top.is_children(node):
                self.pop_from_stack()
            self.push_to_stack(node)
        # 出栈最后一个
        self.debug("剩余: {}".format(', '.join([n.key for n in self.stack])))

    def get_nodes_by_queryset(self, queryset):
        nodes = []
        for n in queryset:
            node = self.get_node_by_key(n.key)
            if not node:
                continue
            nodes.append(node)
        return nodes

    def get_node_by_key(self, key):
        return self._nodes.get(key)

    def debug(self, msg):
        self._debug and logger.debug(msg)

    def set_assets_amount(self):
        for node in self._nodes.values():
            node.assets_amount = node._assets_amount

    def set_full_value(self):
        for node in self._nodes.values():
            node.full_value = node._full_value

    @property
    def nodes(self):
        return list(self._nodes.values())

    # 使用给定节点生成一颗树
    # 找到他们的祖先节点
    # 可选找到他们的子孙节点
    def get_family(self, nodes, with_children=False):
        tree_nodes = set()
        for n in nodes:
            node = self.get_node_by_key(n.key)
            if not node:
                continue
            tree_nodes.update(node._parents)
            tree_nodes.add(node)
            if with_children:
                tree_nodes.update(node._children)
        return list(tree_nodes)

    def get_nodes_parents(self, nodes, with_self=True):
        parents = set()
        for n in nodes:
            node = self.get_node_by_key(n.key)
            parents.update(set(node._parents))
            if with_self:
                parents.add(node)
        return parents


def test_node_tree():
    tree = NodeUtil()
    for node in tree._nodes.values():
        print("Check {}".format(node.key))
        children_wanted = node.get_all_children().count()
        children = len(node._children)
        if children != children_wanted:
            print("{} children not equal: {} != {}".format(node.key, children, children_wanted))

        assets_amount_wanted = node.get_all_assets().count()
        if node._assets_amount != assets_amount_wanted:
            print("{} assets amount not equal: {} != {}".format(
                node.key, node._assets_amount, assets_amount_wanted)
            )

        full_value_wanted = node.full_value
        if node._full_value != full_value_wanted:
            print("{} full value not equal: {} != {}".format(
                node.key, node._full_value, full_value_wanted)
            )

        parents_wanted = node.get_ancestor().count()
        parents = len(node._parents)
        if parents != parents_wanted:
            print("{} parents count not equal: {} != {}".format(
                node.key, parents, parents_wanted)
            )