From 69b16e4754c0f3d473607bc5d3d51e729276f136 Mon Sep 17 00:00:00 2001 From: fit2bot <68588906+fit2bot@users.noreply.github.com> Date: Mon, 19 Dec 2022 11:35:50 +0800 Subject: [PATCH] perf: asset type xpack (#9218) Co-authored-by: feng <1304903146@qq.com> --- apps/assets/api/platform.py | 9 +++++++-- apps/assets/const/base.py | 23 +++++++++++++++++++++++ apps/assets/const/cloud.py | 4 ++++ apps/assets/const/database.py | 6 +++++- apps/assets/const/device.py | 4 ++++ apps/assets/const/host.py | 10 ++++++++-- apps/assets/const/types.py | 20 ++++++++++++++++---- apps/assets/const/web.py | 4 ++++ 8 files changed, 71 insertions(+), 9 deletions(-) diff --git a/apps/assets/api/platform.py b/apps/assets/api/platform.py index 411bf8b86..dbfcc2e4c 100644 --- a/apps/assets/api/platform.py +++ b/apps/assets/api/platform.py @@ -1,10 +1,10 @@ - +from jumpserver.utils import has_valid_xpack_license from common.drf.api import JMSModelViewSet from common.drf.serializers import GroupedChoiceSerializer from assets.models import Platform +from assets.const import AllTypes from assets.serializers import PlatformSerializer - __all__ = ['AssetPlatformViewSet'] @@ -22,6 +22,11 @@ class AssetPlatformViewSet(JMSModelViewSet): 'ops_methods': 'assets.view_platform' } + def get_queryset(self): + queryset = super().get_queryset() + queryset = queryset.filter(type__in=AllTypes.get_types()) + return queryset + def get_object(self): pk = self.kwargs.get('pk', '') if pk.isnumeric(): diff --git a/apps/assets/const/base.py b/apps/assets/const/base.py index 35ad76757..aa5070d5e 100644 --- a/apps/assets/const/base.py +++ b/apps/assets/const/base.py @@ -1,5 +1,6 @@ from django.db.models import TextChoices +from jumpserver.utils import has_valid_xpack_license from .protocol import Protocol @@ -53,3 +54,25 @@ class BaseType(TextChoices): @classmethod def internal_platforms(cls): raise NotImplementedError + + @classmethod + def get_community_types(cls): + raise NotImplementedError + + @classmethod + def get_types(cls): + tps = [tp for tp in cls] + if not has_valid_xpack_license(): + tps = cls.get_community_types() + return tps + + @classmethod + def get_choices(cls): + tps = cls.get_types() + cls_choices = cls.choices + return [ + choice for choice in cls_choices + if choice[0] in tps + ] + + diff --git a/apps/assets/const/cloud.py b/apps/assets/const/cloud.py index 7bc1864f1..1e16c9b13 100644 --- a/apps/assets/const/cloud.py +++ b/apps/assets/const/cloud.py @@ -49,3 +49,7 @@ class CloudTypes(BaseType): cls.PRIVATE: [{'name': 'Vmware-vSphere'}], cls.K8S: [{'name': 'Kubernetes'}], } + + @classmethod + def get_community_types(cls): + return [cls.K8S] diff --git a/apps/assets/const/database.py b/apps/assets/const/database.py index 40b9f8aff..7df58e0a5 100644 --- a/apps/assets/const/database.py +++ b/apps/assets/const/database.py @@ -1,4 +1,3 @@ - from .base import BaseType @@ -62,3 +61,8 @@ class DatabaseTypes(BaseType): cls.REDIS: [{'name': 'Redis'}], } + @classmethod + def get_community_types(cls): + return [ + cls.MYSQL, cls.MARIADB, cls.MONGODB, cls.REDIS + ] diff --git a/apps/assets/const/device.py b/apps/assets/const/device.py index 1e2a5b717..45cfb05e7 100644 --- a/apps/assets/const/device.py +++ b/apps/assets/const/device.py @@ -52,3 +52,7 @@ class DeviceTypes(BaseType): cls.ROUTER: [], cls.FIREWALL: [] } + + @classmethod + def get_community_types(cls): + return [] diff --git a/apps/assets/const/host.py b/apps/assets/const/host.py index 8be44db6f..ed0f9a631 100644 --- a/apps/assets/const/host.py +++ b/apps/assets/const/host.py @@ -34,7 +34,7 @@ class HostTypes(BaseType): def _get_protocol_constrains(cls) -> dict: return { '*': { - 'choices': ['ssh', 'telnet', 'vnc', 'rdp'] + 'choices': ['ssh', 'telnet', 'vnc', 'rdp'] }, cls.WINDOWS: { 'choices': ['rdp', 'ssh', 'vnc'] @@ -97,7 +97,7 @@ class HostTypes(BaseType): { 'name': 'RemoteAppHost', '_protocols': ['rdp', 'ssh'], - 'protocols_setting': { + 'protocols_setting': { 'ssh': { 'required': True } @@ -106,3 +106,9 @@ class HostTypes(BaseType): ], cls.OTHER_HOST: [] } + + @classmethod + def get_community_types(cls) -> list: + return [ + cls.LINUX, cls.UNIX, cls.WINDOWS, cls.OTHER_HOST + ] diff --git a/apps/assets/const/types.py b/apps/assets/const/types.py index c9012cd8c..6a0881e5d 100644 --- a/apps/assets/const/types.py +++ b/apps/assets/const/types.py @@ -62,14 +62,18 @@ class AllTypes(ChoicesMixin): @classmethod def types(cls, with_constraints=True): types = [] - for category, tps in cls.category_types(): + for category, type_cls in cls.category_types(): + tps = type_cls.get_types() types.extend([cls.serialize_type(category, tp, with_constraints) for tp in tps]) return types @classmethod def categories(cls, with_constraints=True): categories = [] - for category, tps in cls.category_types(): + for category, type_cls in cls.category_types(): + tps = type_cls.get_types() + if not tps: + continue category_data = { 'value': category.value, 'label': category.label, @@ -121,6 +125,13 @@ class AllTypes(ChoicesMixin): (Category.CLOUD, CloudTypes) ) + @classmethod + def get_types(cls): + tps = [] + for i in dict(cls.category_types()).values(): + tps.extend(i.get_types()) + return tps + @staticmethod def choice_to_node(choice, pid, opened=True, is_parent=True, meta=None): node = TreeNode(**{ @@ -168,14 +179,15 @@ class AllTypes(ChoicesMixin): root = TreeNode(id='ROOT', name='所有类型', title='所有类型', open=True, isParent=True) nodes = [root] - for category, types in cls.category_types(): + for category, type_cls in cls.category_types(): meta = {'type': 'category', 'category': category.value} category_node = cls.choice_to_node(category, 'ROOT', meta=meta) category_count = category_type_mapper.get(category, 0) category_node.name += f'({category_count})' nodes.append(category_node) - for tp in types: + tps = type_cls.get_types() + for tp in tps: meta = {'type': 'type', 'category': category.value, '_type': tp.value} tp_node = cls.choice_to_node(tp, category_node.id, opened=False, meta=meta) tp_count = category_type_mapper.get(category + '_' + tp, 0) diff --git a/apps/assets/const/web.py b/apps/assets/const/web.py index 20c35b3a1..061030ab9 100644 --- a/apps/assets/const/web.py +++ b/apps/assets/const/web.py @@ -44,3 +44,7 @@ class WebTypes(BaseType): {'name': 'Website'}, ], } + + @classmethod + def get_community_types(cls): + return []