mirror of
https://github.com/jumpserver/jumpserver.git
synced 2026-03-19 11:32:09 +00:00
132 lines
4.5 KiB
Python
132 lines
4.5 KiB
Python
from rest_framework import generics
|
|
from rest_framework.response import Response
|
|
|
|
from rbac.permissions import RBACPermission
|
|
from .mixin import SerializeToTreeNodeMixin
|
|
from assets.models import Node
|
|
from assets.tree.asset_tree import AssetGenericTree
|
|
from assets.tree.node_tree import AssetNodeTree
|
|
from assets.tree.type_tree import AssetTypeTree
|
|
from orgs.utils import current_org, tmp_to_org
|
|
from orgs.models import Organization
|
|
from common.utils import timeit
|
|
|
|
|
|
__all__ = ['AssetGenericTreeApi', 'AssetNodeTreeApi', 'AssetTypeTreeApi']
|
|
|
|
|
|
class AssetGenericTreeApi(SerializeToTreeNodeMixin, generics.ListAPIView):
|
|
permission_classes = (RBACPermission, )
|
|
with_assets = False
|
|
expand_node_key = None
|
|
with_assets_amount = False
|
|
|
|
rbac_perms = {
|
|
'list': 'assets.view_asset',
|
|
'GET': 'assets.view_asset',
|
|
'OPTIONS': 'assets.view_asset',
|
|
}
|
|
def initial(self, request, *args, **kwargs):
|
|
super().initial(request, *args, **kwargs)
|
|
self.initial_node_root_if_need()
|
|
if self.request.query_params.get('assets') == '1':
|
|
self.with_assets = True
|
|
if key := self.request.query_params.get('key'):
|
|
self.expand_node_key = key
|
|
if self.request.query_params.get('asset_amount') == '1':
|
|
self.with_assets_amount = True
|
|
|
|
@property
|
|
def tree_user(self):
|
|
return self.request.user
|
|
|
|
@property
|
|
def tree_orgs(self):
|
|
if current_org.is_root():
|
|
orgs = self.tree_user.orgs.all()
|
|
else:
|
|
orgs = Organization.objects.filter(id=current_org.id)
|
|
return orgs
|
|
|
|
def initial_node_root_if_need(self):
|
|
for org in self.tree_orgs:
|
|
with tmp_to_org(org):
|
|
Node.org_root()
|
|
|
|
@timeit
|
|
def list(self, request, *args, **kwargs):
|
|
data = []
|
|
for org in self.tree_orgs:
|
|
tree = self.get_asset_tree(org)
|
|
tree_data = self.get_tree_data(tree)
|
|
data.extend(tree_data)
|
|
return Response(data=data)
|
|
|
|
def get_asset_tree(self, org) -> AssetGenericTree:
|
|
raise NotImplementedError()
|
|
|
|
def get_tree_data(self, tree):
|
|
tree: AssetGenericTree
|
|
if tree.empty():
|
|
return []
|
|
|
|
with_assets = self.with_assets
|
|
with_assets_amount = self.with_assets_amount
|
|
expand_node_key = self.expand_node_key
|
|
search = self.request.query_params.get('search')
|
|
if search:
|
|
getattr(tree, 'set_use_cache', lambda: None)()
|
|
nodes = tree.filter_nodes(keyword=search)
|
|
ancestors = tree.get_ancestors_of_nodes(nodes)
|
|
nodes = tree.merge_nodes(ancestors, nodes)
|
|
data_nodes = self.serialize_nodes(ancestors, with_asset_amount=with_assets_amount)
|
|
if with_assets:
|
|
assets = tree.filter_assets(keyword=search)
|
|
node_key = tree.root.key
|
|
data_assets = self.serialize_assets(assets, node_key=node_key)
|
|
return data_nodes + data_assets
|
|
return data_nodes
|
|
|
|
if expand_node_key:
|
|
getattr(tree, 'set_use_cache', lambda: None)()
|
|
node = tree.get_node(expand_node_key)
|
|
if not node:
|
|
return []
|
|
nodes = node.children
|
|
data_nodes = self.serialize_nodes(nodes, with_asset_amount=with_assets_amount)
|
|
if with_assets:
|
|
assets = tree.filter_assets(node_key=expand_node_key)
|
|
data_assets = self.serialize_assets(assets, node_key=expand_node_key)
|
|
return data_nodes + data_assets
|
|
return data_nodes
|
|
|
|
data_assets = []
|
|
if current_org.is_root():
|
|
nodes = [tree.root]
|
|
else:
|
|
nodes = [tree.root] + tree.root.children
|
|
if with_assets:
|
|
assets = tree.filter_assets(node_key=tree.root.key)
|
|
data_assets = self.serialize_assets(assets, node_key=tree.root.key)
|
|
extra_nodes = getattr(tree, 'get_extra_nodes', lambda: [])()
|
|
nodes = extra_nodes + nodes
|
|
data_nodes = self.serialize_nodes(nodes, with_asset_amount=with_assets_amount)
|
|
data = data_nodes + data_assets
|
|
return data
|
|
|
|
|
|
class AssetNodeTreeApi(AssetGenericTreeApi):
|
|
|
|
def get_asset_tree(self, org):
|
|
category = self.request.query_params.get('category')
|
|
tree = AssetNodeTree(category=category, org=org)
|
|
tree.init()
|
|
return tree
|
|
|
|
|
|
class AssetTypeTreeApi(AssetGenericTreeApi):
|
|
|
|
def get_asset_tree(self, org):
|
|
tree = AssetTypeTree(org=org)
|
|
tree.init()
|
|
return tree |