diff --git a/apps/assets/const.py b/apps/assets/const.py index d5bf84727..f2ec1a882 100644 --- a/apps/assets/const.py +++ b/apps/assets/const.py @@ -1,6 +1,6 @@ from django.db import models from django.utils.translation import gettext_lazy as _ -from common.db.models import IncludesTextChoicesMeta +from common.db.models import IncludesTextChoicesMeta, ChoicesMixin from common.tree import TreeNode @@ -24,7 +24,7 @@ class PlatformMixin: } -class Category(PlatformMixin, models.TextChoices): +class Category(PlatformMixin, ChoicesMixin, models.TextChoices): HOST = 'host', _('Host') NETWORK = 'network', _("NetworkDevice") DATABASE = 'database', _("Database") @@ -60,7 +60,7 @@ class Category(PlatformMixin, models.TextChoices): } -class HostTypes(PlatformMixin, models.TextChoices): +class HostTypes(PlatformMixin, ChoicesMixin, models.TextChoices): LINUX = 'linux', 'Linux' WINDOWS = 'windows', 'Windows' UNIX = 'unix', 'Unix' @@ -84,14 +84,14 @@ class HostTypes(PlatformMixin, models.TextChoices): } -class NetworkTypes(PlatformMixin, models.TextChoices): +class NetworkTypes(PlatformMixin, ChoicesMixin, models.TextChoices): SWITCH = 'switch', _("Switch") ROUTER = 'router', _("Router") FIREWALL = 'firewall', _("Firewall") OTHER_NETWORK = 'other_network', _("Other device") -class DatabaseTypes(PlatformMixin, models.TextChoices): +class DatabaseTypes(PlatformMixin, ChoicesMixin, models.TextChoices): MYSQL = 'mysql', 'MySQL' MARIADB = 'mariadb', 'MariaDB' POSTGRESQL = 'postgresql', 'PostgreSQL' @@ -110,15 +110,15 @@ class DatabaseTypes(PlatformMixin, models.TextChoices): return meta -class WebTypes(PlatformMixin, models.TextChoices): +class WebTypes(PlatformMixin, ChoicesMixin, models.TextChoices): General = 'general', 'General' -class CloudTypes(PlatformMixin, models.TextChoices): +class CloudTypes(PlatformMixin, ChoicesMixin, models.TextChoices): K8S = 'k8s', 'Kubernetes' -class AllTypes(metaclass=IncludesTextChoicesMeta): +class AllTypes(ChoicesMixin, metaclass=IncludesTextChoicesMeta): choices: list includes = [ HostTypes, NetworkTypes, DatabaseTypes, @@ -202,7 +202,7 @@ class AllTypes(metaclass=IncludesTextChoicesMeta): return nodes -class Protocol(models.TextChoices): +class Protocol(ChoicesMixin, models.TextChoices): ssh = 'ssh', 'SSH' rdp = 'rdp', 'RDP' telnet = 'telnet', 'Telnet' diff --git a/apps/assets/models/asset/common.py b/apps/assets/models/asset/common.py index 1e4c3d71f..b4b4e5b81 100644 --- a/apps/assets/models/asset/common.py +++ b/apps/assets/models/asset/common.py @@ -10,6 +10,7 @@ from django.db import models from django.utils.translation import ugettext_lazy as _ from orgs.mixins.models import OrgManager, JMSOrgBaseModel +from ...const import Category from ..platform import Platform from ..base import AbsConnectivity @@ -121,7 +122,8 @@ class Asset(AbsConnectivity, NodesRelationMixin, JMSOrgBaseModel): @property def type_display(self): - return self.platform.type + value = self.platform.type + return Category.get_label(value) @property def category(self): @@ -129,7 +131,8 @@ class Asset(AbsConnectivity, NodesRelationMixin, JMSOrgBaseModel): @property def category_display(self): - return self.platform.category + value = self.platform.category + return Category.get_label(value) def as_node(self): from assets.models import Node diff --git a/apps/common/db/models.py b/apps/common/db/models.py index 03e95ca5d..f180b0da1 100644 --- a/apps/common/db/models.py +++ b/apps/common/db/models.py @@ -85,6 +85,14 @@ class BitOperationChoice: return [(cls.NAME_MAP[i], j) for i, j in cls.DB_CHOICES] +class ChoicesMixin: + _value2label_map_: dict + + @classmethod + def get_label(cls, value: (str, int)): + return cls._value2label_map_[value] + + class BaseCreateUpdateModel(models.Model): created_by = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Created by')) updated_by = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Updated by'))