jumpserver/apps/perms/api/user_permission/tree/node_with_asset.py

import abc
from urllib.parse import parse_qsl

from django.conf import settings
from django.db.models import F, Value, CharField
from rest_framework.exceptions import PermissionDenied, NotFound
from rest_framework.generics import ListAPIView
from rest_framework.generics import get_object_or_404
from rest_framework.request import Request
from rest_framework.response import Response

from accounts.const import AliasAccount
from assets.api import SerializeToTreeNodeMixin
from assets.models import Asset
from assets.utils import KubernetesTree
from authentication.models import ConnectionToken
from common.utils import get_object_or_none, lazyproperty
from common.utils.common import timeit
from perms.hands import Node
from perms.models import PermNode
from perms.utils import PermAssetDetailUtil, UserPermNodeUtil
from perms.utils import UserPermAssetUtil
from .mixin import RebuildTreeMixin
from ..mixin import SelfOrPKUserMixin

__all__ = [
    'UserGrantedK8sAsTreeApi',
    'UserPermedNodesWithAssetsAsTreeApi',
    'UserPermedNodeChildrenWithAssetsAsTreeApi',
    'UserPermedNodeChildrenWithAssetsAsCategoryTreeApi',
]


class BaseUserNodeWithAssetAsTreeApi(
    SelfOrPKUserMixin, RebuildTreeMixin,
    SerializeToTreeNodeMixin, ListAPIView
):
    page_limit = 10000

    def list(self, request, *args, **kwargs):
        offset = int(request.query_params.get('offset', 0))
        page_assets = self.get_page_assets()

        if not offset:
            nodes, assets = self.get_nodes_assets()
            page = page_assets[:self.page_limit]
            assets = [*assets, *page]
            tree_nodes = self.serialize_nodes(nodes, with_asset_amount=True)
            tree_assets = self.serialize_assets(assets, **self.serialize_asset_kwargs)
            data = list(tree_nodes) + list(tree_assets)
        else:
            page = page_assets[offset:(offset + self.page_limit)]
            data = self.serialize_assets(page, **self.serialize_asset_kwargs) if page else []
        offset += len(page)
        headers = {'X-JMS-TREE-OFFSET': offset} if offset else {}
        return Response(data=data, headers=headers)

    @abc.abstractmethod
    def get_nodes_assets(self):
        return [], []

    def get_page_assets(self):
        return []

    @property
    def serialize_asset_kwargs(self):
        return {}


class UserPermedNodesWithAssetsAsTreeApi(BaseUserNodeWithAssetAsTreeApi):
    query_node_util: UserPermNodeUtil
    query_asset_util: UserPermAssetUtil

    def get_nodes_assets(self):
        self.query_node_util = UserPermNodeUtil(self.request.user)
        ung_nodes, ung_assets = self._get_nodes_assets_for_ungrouped()
        fav_nodes, fav_assets = self._get_nodes_assets_for_favorite()
        all_nodes, all_assets = self._get_nodes_assets_for_all()
        nodes = list(ung_nodes) + list(fav_nodes) + list(all_nodes)
        assets = list(ung_assets) + list(fav_assets) + list(all_assets)
        return nodes, assets

    def get_page_assets(self):
        return self.query_asset_util.get_all_assets().annotate(parent_key=F('nodes__key'))

    @timeit
    def _get_nodes_assets_for_ungrouped(self):
        if not settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
            return [], []
        node = self.query_node_util.get_ungrouped_node()
        assets = self.query_asset_util.get_ungroup_assets()
        assets = assets.annotate(parent_key=Value(node.key, output_field=CharField()))
        return [node], assets

    @lazyproperty
    def query_asset_util(self):
        return UserPermAssetUtil(self.user)

    @timeit
    def _get_nodes_assets_for_favorite(self):
        node = self.query_node_util.get_favorite_node()
        assets = self.query_asset_util.get_favorite_assets()
        assets = assets.annotate(parent_key=Value(node.key, output_field=CharField()))
        return [node], assets

    @timeit
    def _get_nodes_assets_for_all(self):
        nodes = self.query_node_util.get_whole_tree_nodes(with_special=False)
        if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
            assets = self.query_asset_util.get_perm_nodes_assets()
        else:
            assets = Asset.objects.none()
        assets = assets.annotate(parent_key=F('nodes__key'))
        return nodes, assets


class UserPermedNodeChildrenWithAssetsAsTreeApi(BaseUserNodeWithAssetAsTreeApi):
    """ 用户授权的节点的子节点与资产树 """

    # 默认展开的节点key
    default_unfolded_node_key = None

    @timeit
    def get_nodes_assets(self):
        query_node_util = UserPermNodeUtil(self.user)
        query_asset_util = UserPermAssetUtil(self.user)
        node_key = self.query_node_key
        if not node_key:
            nodes, unfolded_node = query_node_util.get_top_level_nodes(with_unfolded_node=True)
            if unfolded_node:
                """ 默认展开的节点, 获取根节点下的资产 """
                assets = query_asset_util.get_node_assets(key=unfolded_node.key)
                self.default_unfolded_node_key = unfolded_node.key
            else:
                assets = Asset.objects.none()
        elif node_key == PermNode.UNGROUPED_NODE_KEY:
            nodes = PermNode.objects.none()
            assets = query_asset_util.get_ungroup_assets()
        elif node_key == PermNode.FAVORITE_NODE_KEY:
            nodes = PermNode.objects.none()
            assets = query_asset_util.get_favorite_assets()
        else:
            nodes = query_node_util.get_node_children(node_key)
            assets = query_asset_util.get_node_assets(key=node_key)
        assets = assets.prefetch_related('platform')
        return nodes, assets

    @lazyproperty
    def query_node_key(self):
        node_key = self.request.query_params.get('key', None)
        if node_key is None:
            node_id = self.request.query_params.get('id', None)
            node = get_object_or_none(Node, id=node_id)
            node_key = getattr(node, 'key', None)
        return node_key

    @property
    def serialize_asset_kwargs(self):
        return {
            'node_key': self.query_node_key or self.default_unfolded_node_key
        }


class UserPermedNodeChildrenWithAssetsAsCategoryTreeApi(BaseUserNodeWithAssetAsTreeApi):
    @property
    def is_sync(self):
        sync = self.request.query_params.get('sync', 0)
        return int(sync) == 1

    @property
    def tp(self):
        params = self.request.query_params
        return [params.get('category'), params.get('type')]

    @lazyproperty
    def query_asset_util(self):
        return UserPermAssetUtil(self.user)

    @timeit
    def get_assets(self):
        return self.query_asset_util.get_all_assets()

    def _get_tree_nodes_async(self):
        if not self.tp or not all(self.tp):
            nodes = UserPermAssetUtil.get_type_nodes_tree_or_cached(self.user)
            return nodes, []

        category, tp = self.tp
        assets = self.get_assets().filter(platform__type=tp, platform__category=category)
        return [], assets

    def _get_tree_nodes_sync(self):
        if self.request.query_params.get('lv'):
            return []
        nodes = self.query_asset_util.get_type_nodes_tree()
        return nodes, []

    @property
    def serialize_asset_kwargs(self):
        return {
            'get_pid': lambda asset, platform: 'ROOT_{}_{}'.format(platform.category.upper(), platform.type),
        }

    def serialize_nodes(self, nodes, with_asset_amount=False):
        return nodes

    def get_nodes_assets(self):
        if self.is_sync:
            return self._get_tree_nodes_sync()
        else:
            return self._get_tree_nodes_async()

    def get_page_assets(self):
        if self.is_sync:
            return self.get_assets()
        else:
            return []


class UserGrantedK8sAsTreeApi(SelfOrPKUserMixin, ListAPIView):
    """ 用户授权的K8s树 """

    def get_token(self):
        token_id = self.request.query_params.get('token')
        token = get_object_or_404(ConnectionToken, pk=token_id)
        if token.is_expired:
            raise PermissionDenied('Token is expired')
        token.renewal()
        return token

    def get_account_secret(self, token: ConnectionToken):
        util = PermAssetDetailUtil(self.user, token.asset)
        accounts = util.get_permed_accounts_for_user()
        account_name = token.account

        if account_name in [AliasAccount.INPUT, AliasAccount.USER]:
            return token.input_secret
        else:
            accounts = filter(lambda x: x.name == account_name, accounts)
            accounts = list(accounts)
            if not accounts:
                raise NotFound('Account is not found')
            account = accounts[0]
            return account.secret

    @staticmethod
    def get_namespace_and_pod(key):
        namespace_and_pod = dict(parse_qsl(key))
        pod = namespace_and_pod.get('pod')
        namespace = namespace_and_pod.get('namespace')
        return namespace, pod

    def list(self, request: Request, *args, **kwargs):
        token = self.get_token()
        asset = token.asset
        secret = self.get_account_secret(token)
        key = self.request.query_params.get('key')
        namespace, pod = self.get_namespace_and_pod(key)

        tree = []
        k8s_tree_instance = KubernetesTree(asset, secret)
        if not any([namespace, pod]) and not key:
            asset_node = k8s_tree_instance.as_asset_tree_node()
            tree.append(asset_node)
        tree.extend(k8s_tree_instance.async_tree_node(namespace, pod))
        return Response(data=tree)