perf: 修改 base

apps/acls/serializers/login_asset_acl.py

@@ -1,9 +1,11 @@
 from rest_framework import serializers
 from django.utils.translation import ugettext_lazy as _
+
 from orgs.mixins.serializers import BulkOrgResourceModelSerializer
-from assets.models import SystemUser
-from acls import models
 from orgs.models import Organization
+from assets.models import SystemUser
+from assets.const import Protocol
+from acls import models
 
 
 __all__ = ['LoginAssetACLSerializer']
@@ -54,7 +56,7 @@ class LoginAssetACLSystemUsersSerializer(serializers.Serializer):
     protocol_group = serializers.ListField(
         default=['*'], child=serializers.CharField(max_length=16), label=_('Protocol'),
         help_text=protocol_group_help_text.format(
-            ', '.join([SystemUser.Protocol.ssh, SystemUser.Protocol.telnet])
+            ', '.join([Protocol.ssh, Protocol.telnet])
         )
     )
 

apps/assets/api/system_user.py

@@ -10,6 +10,7 @@ from common.mixins.api import SuggestionMixin
 from orgs.mixins.api import OrgBulkModelViewSet
 from orgs.mixins import generics
 from orgs.utils import tmp_to_root_org
+from assets.const import Protocol
 from ..models import SystemUser, CommandFilterRule
 from .. import serializers
 from ..serializers import SystemUserWithAuthInfoSerializer, SystemUserTempAuthSerializer
@@ -56,7 +57,7 @@ class SystemUserViewSet(SuggestionMixin, OrgBulkModelViewSet):
         """ API 获取可选的 su_from 系统用户"""
         queryset = self.filter_queryset(self.get_queryset())
         queryset = queryset.filter(
-            protocol=SystemUser.Protocol.ssh, login_mode=SystemUser.LOGIN_AUTO
+            protocol=Protocol.ssh, login_mode=SystemUser.LOGIN_AUTO
         )
         return self.get_paginate_response_if_need(queryset)
 

apps/assets/const.py

@@ -1,2 +1,92 @@
-# -*- coding: utf-8 -*-
+from django.db import models
-#
+from django.utils.translation import gettext_lazy as _
+from common.db.models import IncludesTextChoicesMeta
+
+
+__all__ = [
+    'Category', 'HostTypes', 'NetworkTypes', 'DatabaseTypes',
+    'RemoteAppTypes', 'CloudTypes', 'Protocol', 'AllTypes',
+]
+
+
+class Category(models.TextChoices):
+    HOST = 'host', _('Host')
+    NETWORK = 'network', _("Networking")
+    DATABASE = 'database', _("Database")
+    REMOTE_APP = 'remote_app', _("Remote app")
+    CLOUD = 'cloud', _("Clouding")
+
+
+class HostTypes(models.TextChoices):
+    LINUX = 'linux', 'Linux'
+    WINDOWS = 'windows', 'Windows'
+    UNIX = 'unix', 'Unix'
+    BSD = 'bsd', 'BSD'
+    MACOS = 'macos', 'MacOS'
+    MAINFRAME = 'mainframe', _("Mainframe")
+    OTHER_HOST = 'other_host', _("Other host")
+
+
+class NetworkTypes(models.TextChoices):
+    SWITCH = 'switch', _("Switch")
+    ROUTER = 'router', _("Router")
+    FIREWALL = 'firewall', _("Firewall")
+    OTHER_NETWORK = 'other_network', _("Other device")
+
+
+class DatabaseTypes(models.TextChoices):
+    MYSQL = 'mysql', 'MySQL'
+    MARIADB = 'mariadb', 'MariaDB'
+    POSTGRESQL = 'postgresql', 'PostgreSQL'
+    ORACLE = 'oracle', 'Oracle'
+    SQLSERVER = 'sqlserver', 'SQLServer'
+    MONGODB = 'mongodb', 'MongoDB'
+    REDIS = 'redis', 'Redis'
+
+
+class RemoteAppTypes(models.TextChoices):
+    CHROME = 'chrome', 'Chrome'
+    VSPHERE = 'vsphere', 'vSphere client'
+    MYSQL_WORKBENCH = 'mysql_workbench', 'MySQL workbench'
+    CUSTOM_REMOTE_APP = 'custom_remote_app', _("Custom")
+
+
+class CloudTypes(models.TextChoices):
+    K8S = 'k8s', 'Kubernetes'
+
+
+class AllTypes(metaclass=IncludesTextChoicesMeta):
+    choices: list
+    includes = [
+        HostTypes, NetworkTypes, DatabaseTypes,
+        RemoteAppTypes, CloudTypes
+    ]
+
+
+class Protocol(models.TextChoices):
+    ssh = 'ssh', 'SSH'
+    rdp = 'rdp', 'RDP'
+    telnet = 'telnet', 'Telnet'
+    vnc = 'vnc', 'VNC'
+
+    mysql = 'mysql', 'MySQL'
+    mariadb = 'mariadb', 'MariaDB'
+    oracle = 'oracle', 'Oracle'
+    postgresql = 'postgresql', 'PostgreSQL'
+    sqlserver = 'sqlserver', 'SQLServer'
+    redis = 'redis', 'Redis'
+    mongodb = 'mongodb', 'MongoDB'
+
+    k8s = 'k8s', 'K8S'
+
+    @classmethod
+    def host_protocols(cls):
+        return [cls.ssh, cls.rdp, cls.telnet, cls.vnc]
+
+    @classmethod
+    def db_protocols(cls):
+        return [
+            cls.mysql, cls.mariadb, cls.postgresql, cls.oracle,
+            cls.sqlserver, cls.redis, cls.mongodb,
+        ]
+

apps/assets/migrations/0097_auto_20220407_1726.py

@@ -0,0 +1,34 @@
+# Generated by Django 3.1.14 on 2022-04-07 09:26
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('assets', '0096_auto_20220406_1546'),
+    ]
+
+    operations = [
+        migrations.RenameField(
+            model_name='platform',
+            old_name='base',
+            new_name='type',
+        ),
+        migrations.AddField(
+            model_name='platform',
+            name='category',
+            field=models.CharField(choices=[('host', 'Host'), ('network', 'Networking'), ('database', 'Database'), ('remote_app', 'Remote app'), ('cloud', 'Clouding')], default='host', max_length=16, verbose_name='Category'),
+            preserve_default=False,
+        ),
+        migrations.AlterField(
+            model_name='asset',
+            name='type',
+            field=models.CharField(choices=[('linux', 'Linux'), ('unix', 'Unix'), ('windows', 'Windows'), ('macos', 'MacOS'), ('mainframe', 'Mainframe'), ('other_host', 'Other host'), ('switch', 'Switch'), ('router', 'Router'), ('firewall', 'Firewall'), ('other_network', 'Other device'), ('mysql', 'MySQL'), ('mariadb', 'MariaDB'), ('postgresql', 'PostgreSQL'), ('oracle', 'Oracle'), ('sqlserver', 'SQLServer'), ('mongodb', 'MongoDB'), ('redis', 'Redis'), ('chrome', 'Chrome'), ('vsphere', 'vSphere client'), ('mysql_workbench', 'MySQL workbench'), ('custom_remote_app', 'Custom'), ('k8s', 'Kubernetes')], max_length=128, verbose_name='Type'),
+        ),
+        migrations.AlterField(
+            model_name='systemuser',
+            name='protocol',
+            field=models.CharField(choices=[('ssh', 'SSH'), ('rdp', 'RDP'), ('telnet', 'Telnet'), ('vnc', 'VNC'), ('mysql', 'MySQL'), ('mariadb', 'MariaDB'), ('oracle', 'Oracle'), ('postgresql', 'PostgreSQL'), ('sqlserver', 'SQLServer'), ('redis', 'Redis'), ('mongodb', 'MongoDB'), ('k8s', 'K8S')], default='ssh', max_length=16, verbose_name='Protocol'),
+        ),
+    ]

apps/assets/migrations/0098_auto_20220407_1730.py

@@ -0,0 +1,23 @@
+# Generated by Django 3.1.14 on 2022-04-07 09:30
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('assets', '0097_auto_20220407_1726'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='asset',
+            name='type',
+            field=models.CharField(choices=[('linux', 'Linux'), ('windows', 'Windows'), ('unix', 'Unix'), ('bsd', 'BSD'), ('macos', 'MacOS'), ('mainframe', 'Mainframe'), ('other_host', 'Other host'), ('switch', 'Switch'), ('router', 'Router'), ('firewall', 'Firewall'), ('other_network', 'Other device'), ('mysql', 'MySQL'), ('mariadb', 'MariaDB'), ('postgresql', 'PostgreSQL'), ('oracle', 'Oracle'), ('sqlserver', 'SQLServer'), ('mongodb', 'MongoDB'), ('redis', 'Redis'), ('chrome', 'Chrome'), ('vsphere', 'vSphere client'), ('mysql_workbench', 'MySQL workbench'), ('custom_remote_app', 'Custom'), ('k8s', 'Kubernetes')], max_length=128, verbose_name='Type'),
+        ),
+        migrations.AlterField(
+            model_name='platform',
+            name='type',
+            field=models.CharField(choices=[('linux', 'Linux'), ('windows', 'Windows'), ('unix', 'Unix'), ('bsd', 'BSD'), ('macos', 'MacOS'), ('mainframe', 'Mainframe'), ('other_host', 'Other host'), ('switch', 'Switch'), ('router', 'Router'), ('firewall', 'Firewall'), ('other_network', 'Other device'), ('mysql', 'MySQL'), ('mariadb', 'MariaDB'), ('postgresql', 'PostgreSQL'), ('oracle', 'Oracle'), ('sqlserver', 'SQLServer'), ('mongodb', 'MongoDB'), ('redis', 'Redis'), ('chrome', 'Chrome'), ('vsphere', 'vSphere client'), ('mysql_workbench', 'MySQL workbench'), ('custom_remote_app', 'Custom'), ('k8s', 'Kubernetes')], default='Linux', max_length=32, verbose_name='Base'),
+        ),
+    ]

apps/assets/models/asset/_category.py

@@ -1,73 +0,0 @@
-from django.db import models
-from django.utils.translation import gettext_lazy as _
-
-
-class Category(models.TextChoices):
-    HOST = 'host', _('Host')
-    NETWORK = 'network', _("Networking")
-    DATABASE = 'database', _("Database")
-    REMOTE_APP = 'remote_app', _("Remote app")
-    CLOUD = 'cloud', _("Clouding")
-
-
-class HostTypes(models.TextChoices):
-    LINUX = 'linux', 'Linux'
-    UNIX = 'unix', 'Unix'
-    WINDOWS = 'windows', 'Windows'
-    MACOS = 'macos', 'MacOS'
-    MAINFRAME = 'mainframe', _("Mainframe")
-    OTHER_HOST = 'other_host', _("Other host")
-
-    def __new__(cls, value):
-        """
-        添加 Category
-        :param value:
-        """
-        obj = str.__new__(cls)
-        obj.category = Category.HOST
-        return obj
-
-
-class NetworkTypes(models.TextChoices):
-    SWITCH = 'switch', _("Switch")
-    ROUTER = 'router', _("Router")
-    FIREWALL = 'firewall', _("Firewall")
-    OTHER_NETWORK = 'other_network', _("Other device")
-
-
-class DatabaseTypes(models.TextChoices):
-    MYSQL = 'mysql', 'MySQL'
-    MARIADB = 'mariadb', 'MariaDB'
-    POSTGRESQL = 'postgresql', 'PostgreSQL'
-    ORACLE = 'oracle', 'Oracle'
-    SQLSERVER = 'sqlserver', 'SQLServer'
-    MONGODB = 'mongodb', 'MongoDB'
-    REDIS = 'redis', 'Redis'
-
-
-class RemoteAppTypes(models.TextChoices):
-    CHROME = 'chrome', 'Chrome'
-    VSPHERE = 'vsphere', 'vSphere client'
-    MYSQL_WORKBENCH = 'mysql_workbench', 'MySQL workbench'
-    CUSTOM_REMOTE_APP = 'custom_remote_app', _("Custom")
-
-
-class CloudTypes(models.TextChoices):
-    K8S = 'k8s', 'Kubernetes'
-
-
-class AllTypes:
-    includes = [
-        HostTypes, NetworkTypes, DatabaseTypes,
-        RemoteAppTypes, CloudTypes
-    ]
-
-    @classmethod
-    def choices(cls):
-        choices = []
-        for tp in cls.includes:
-            choices.extend(tp.choices)
-        return choices
-
-
-

apps/assets/models/asset/common.py

@@ -13,9 +13,9 @@ from rest_framework.exceptions import ValidationError
 
 from common.utils import lazyproperty
 from orgs.mixins.models import OrgModelMixin, OrgManager
+from assets.const import Category, AllTypes
 from ..platform import Platform
 from ..base import AbsConnectivity
-from ._category import Category, AllTypes
 
 __all__ = ['Asset', 'ProtocolsMixin', 'AssetQuerySet', 'default_node', 'default_cluster']
 logger = logging.getLogger(__name__)
@@ -123,11 +123,12 @@ class NodesRelationMixin:
 
 
 class Asset(AbsConnectivity, ProtocolsMixin, NodesRelationMixin, OrgModelMixin):
+    Category = Category
     id = models.UUIDField(default=uuid.uuid4, primary_key=True)
     hostname = models.CharField(max_length=128, verbose_name=_('Hostname'))
     ip = models.CharField(max_length=128, verbose_name=_('IP'), db_index=True)
     category = models.CharField(max_length=16, choices=Category.choices, verbose_name=_("Category"))
-    type = models.CharField(max_length=128, choices=AllTypes.choices(), verbose_name=_("Type"))
+    type = models.CharField(max_length=128, choices=AllTypes.choices, verbose_name=_("Type"))
     protocol = models.CharField(max_length=128, default=ProtocolsMixin.Protocol.ssh,
                                 choices=ProtocolsMixin.Protocol.choices, verbose_name=_('Protocol'))
     port = models.IntegerField(default=22, verbose_name=_('Port'))
@@ -183,14 +184,6 @@ class Asset(AbsConnectivity, ProtocolsMixin, NodesRelationMixin, OrgModelMixin):
             return False, warning
         return True, warning
 
-    @property
-    def category_display(self):
-        return self.get_category_display()
-
-    @property
-    def type_display(self):
-        pass
-
     @lazyproperty
     def platform_base(self):
         return self.platform.base

apps/assets/models/asset/database.py

@@ -6,4 +6,3 @@ from .common import Asset
 
 class Database(Asset):
     database = models.CharField(max_length=1024, verbose_name=_("Database"), blank=True)
-

apps/assets/models/asset/host.py

@@ -1,12 +1,15 @@
 from django.utils.translation import gettext_lazy as _
 from django.db import models
 
+from assets.const import Category
 from common.mixins.models import CommonModelMixin
 from .common import Asset
 
 
 class Host(Asset):
-    pass
+    def save(self, *args, **kwargs):
+        self.category = Category.HOST
+        return super().save(*args, **kwargs)
 
 
 class DeviceInfo(CommonModelMixin):

apps/assets/models/asset/network.py

@@ -0,0 +1,5 @@
+from .common import Asset
+
+
+class Network(Asset):
+    pass

apps/assets/models/platform.py

@@ -1,19 +1,12 @@
 from django.db import models
 from django.utils.translation import gettext_lazy as _
 
+from assets.const import Category, AllTypes
 from common.fields.model import JsonDictTextField
 
 __all__ = ['Platform']
 
 
-class Category(models.TextChoices):
-    Host = 'host', _('Host')
-    Network = 'network', _('Network device')
-    Database = 'database', _('Database')
-    RemoteApp = 'remote_app', _('Microsoft remote app')
-    Cloud = 'cloud', _("Cloud")
-
-
 class Platform(models.Model):
     CHARSET_CHOICES = (
         ('utf8', 'UTF-8'),
@@ -28,7 +21,8 @@ class Platform(models.Model):
         ('Other', 'Other'),
     )
     name = models.SlugField(verbose_name=_("Name"), unique=True, allow_unicode=True)
-    base = models.CharField(choices=BASE_CHOICES, max_length=16, default='Linux', verbose_name=_("Base"))
+    category = models.CharField(max_length=16, choices=Category.choices, verbose_name=_("Category"))
+    type = models.CharField(choices=AllTypes.choices, max_length=32, default='Linux', verbose_name=_("Base"))
     charset = models.CharField(default='utf8', choices=CHARSET_CHOICES, max_length=8, verbose_name=_("Charset"))
     meta = JsonDictTextField(blank=True, null=True, verbose_name=_("Meta"))
     internal = models.BooleanField(default=False, verbose_name=_("Internal"))

apps/assets/models/user.py

@@ -10,6 +10,7 @@ from django.core.validators import MinValueValidator, MaxValueValidator
 from django.core.cache import cache
 
 from common.utils import signer, get_object_or_none
+from assets.const import Protocol
 from .base import BaseUser
 from .asset import Asset
 from .authbook import AuthBook
@@ -21,20 +22,7 @@ logger = logging.getLogger(__name__)
 
 class ProtocolMixin:
     protocol: str
-
+    Protocol = Protocol
-    class Protocol(models.TextChoices):
-        ssh = 'ssh', 'SSH'
-        rdp = 'rdp', 'RDP'
-        telnet = 'telnet', 'Telnet'
-        vnc = 'vnc', 'VNC'
-        mysql = 'mysql', 'MySQL'
-        oracle = 'oracle', 'Oracle'
-        mariadb = 'mariadb', 'MariaDB'
-        postgresql = 'postgresql', 'PostgreSQL'
-        sqlserver = 'sqlserver', 'SQLServer'
-        redis = 'redis', 'Redis'
-        mongodb = 'mongodb', 'MongoDB'
-        k8s = 'k8s', 'K8S'
 
     SUPPORT_PUSH_PROTOCOLS = [Protocol.ssh, Protocol.rdp]
 
@@ -245,7 +233,7 @@ class SystemUser(ProtocolMixin, AuthMixin, BaseUser):
     groups = models.ManyToManyField('users.UserGroup', blank=True, verbose_name=_("User groups"))
     type = models.CharField(max_length=16, choices=Type.choices, default=Type.common, verbose_name=_('Type'))
     priority = models.IntegerField(default=81, verbose_name=_("Priority"), help_text=_("1-100, the lower the value will be match first"), validators=[MinValueValidator(1), MaxValueValidator(100)])
-    protocol = models.CharField(max_length=16, choices=ProtocolMixin.Protocol.choices, default='ssh', verbose_name=_('Protocol'))
+    protocol = models.CharField(max_length=16, choices=Protocol.choices, default='ssh', verbose_name=_('Protocol'))
     auto_push = models.BooleanField(default=True, verbose_name=_('Auto push'))
     sudo = models.TextField(default='/bin/whoami', verbose_name=_('Sudo'))
     shell = models.CharField(max_length=64,  default='/bin/bash', verbose_name=_('Shell'))

apps/assets/serializers/asset/common.py

@@ -70,6 +70,8 @@ class AssetSerializer(BulkOrgResourceModelSerializer):
     labels_display = serializers.ListField(
         child=serializers.CharField(), label=_('Labels name'), required=False, read_only=True
     )
+    category_display = serializers.ReadOnlyField(source='get_category_display', label=_("Category display"))
+    type_display = serializers.ReadOnlyField(source='get_type_display', label=_("Type display"))
 
     """
     资产的数据结构
@@ -78,7 +80,7 @@ class AssetSerializer(BulkOrgResourceModelSerializer):
     class Meta:
         model = Asset
         fields_mini = [
-            'id', 'category', 'category_display', 'type',
+            'id', 'category', 'category_display', 'type', 'type_display',
             'hostname', 'ip', 'platform', 'protocols'
         ]
         fields_small = fields_mini + [

apps/assets/serializers/system_user.py

@@ -6,6 +6,7 @@ from common.mixins.serializers import BulkSerializerMixin
 from common.utils import ssh_pubkey_gen
 from common.validators import alphanumeric_re, alphanumeric_cn_re, alphanumeric_win_re
 from orgs.mixins.serializers import BulkOrgResourceModelSerializer
+from assets.const import Protocol
 from ..models import SystemUser, Asset
 from .utils import validate_password_contains_left_double_curly_bracket
 from .base import AuthSerializerMixin
@@ -107,9 +108,9 @@ class SystemUserSerializer(AuthSerializerMixin, BulkOrgResourceModelSerializer):
     def validate_username(self, username):
         protocol = self.get_initial_value("protocol")
         if username:
-            if protocol == SystemUser.Protocol.telnet:
+            if protocol == Protocol.telnet:
                 regx = alphanumeric_cn_re
-            elif protocol == SystemUser.Protocol.rdp:
+            elif protocol == Protocol.rdp:
                 regx = alphanumeric_win_re
             else:
                 regx = alphanumeric_re
@@ -122,8 +123,8 @@ class SystemUserSerializer(AuthSerializerMixin, BulkOrgResourceModelSerializer):
             return ''
 
         login_mode = self.get_initial_value("login_mode")
-        if login_mode == SystemUser.LOGIN_AUTO and protocol != SystemUser.Protocol.vnc \
+        if login_mode == SystemUser.LOGIN_AUTO and protocol != Protocol.vnc \
-                and protocol != SystemUser.Protocol.redis:
+                and protocol != Protocol.redis:
             msg = _('* Automatic login mode must fill in the username.')
             raise serializers.ValidationError(msg)
         return username
@@ -163,8 +164,8 @@ class SystemUserSerializer(AuthSerializerMixin, BulkOrgResourceModelSerializer):
             error = _('This field is required.')
             raise serializers.ValidationError(error)
         # self: protocol ssh
-        protocol = self.get_initial_value('protocol', default=SystemUser.Protocol.ssh.value)
+        protocol = self.get_initial_value('protocol', default=Protocol.ssh.value)
-        if protocol not in [SystemUser.Protocol.ssh.value]:
+        if protocol not in [Protocol.ssh.value]:
             error = _('Only ssh protocol system users are allowed')
             raise serializers.ValidationError(error)
         # su_from: protocol same
@@ -184,7 +185,7 @@ class SystemUserSerializer(AuthSerializerMixin, BulkOrgResourceModelSerializer):
             tp = attrs.get('type')
         if tp != SystemUser.Type.admin:
             return attrs
-        attrs['protocol'] = SystemUser.Protocol.ssh
+        attrs['protocol'] = Protocol.ssh
         attrs['login_mode'] = SystemUser.LOGIN_AUTO
         attrs['username_same_with_user'] = False
         attrs['auto_push'] = False
@@ -202,7 +203,7 @@ class SystemUserSerializer(AuthSerializerMixin, BulkOrgResourceModelSerializer):
         if auto_gen_key and not self.instance:
             password = SystemUser.gen_password()
             attrs['password'] = password
-            if protocol == SystemUser.Protocol.ssh:
+            if protocol == Protocol.ssh:
                 private_key, public_key = SystemUser.gen_key(username)
                 attrs['private_key'] = private_key
                 attrs['public_key'] = public_key

apps/authentication/models.py

@@ -3,8 +3,9 @@ import uuid
 from django.utils.translation import ugettext_lazy as _
 from rest_framework.authtoken.models import Token
 from django.conf import settings
+from django.db import models
 
-from common.db import models
+from common.db.models import BaseCreateUpdateModel
 
 
 class AccessKey(models.Model):
@@ -40,7 +41,7 @@ class PrivateToken(Token):
         verbose_name = _('Private Token')
 
 
-class SSOToken(models.JMSBaseModel):
+class SSOToken(BaseCreateUpdateModel):
     """
     类似腾讯企业邮的 [单点登录](https://exmail.qq.com/qy_mng_logic/doc#10036)
     出于安全考虑，这里的 `token` 使用一次随即过期。但我们保留每一个生成过的 `token`。
@@ -53,7 +54,7 @@ class SSOToken(models.JMSBaseModel):
         verbose_name = _('SSO token')
 
 
-class ConnectionToken(models.JMSBaseModel):
+class ConnectionToken(BaseCreateUpdateModel):
     # Todo: 未来可能放到这里，不记录到 redis 了，虽然方便，但是不易于审计
     # Todo: add connection token 可能要授权给 普通用户, 或者放开就行
 

apps/common/db/models.py

@@ -13,57 +13,34 @@ import uuid
 from functools import reduce, partial
 import inspect
 
-from django.db.models import *
+from django.db import models
+from django.db.models import F, Value, ExpressionWrapper
+from enum import _EnumDict
 from django.db.models import QuerySet
 from django.db.models.functions import Concat
 from django.utils.translation import ugettext_lazy as _
 
 
-class Choice(str):
+class IncludesTextChoicesMeta(type):
-    def __new__(cls, value, label=''):  # `deepcopy` 的时候不会传 `label`
+    def __new__(metacls, classname, bases, classdict):
-        self = super().__new__(cls, value)
+        includes = classdict.pop('includes', None)
-        self.label = label
+        assert includes
-        return self
 
+        attrs = _EnumDict()
+        for k, v in classdict.items():
+            attrs[k] = v
 
-class ChoiceSetType(type):
+        for cls in includes:
-    def __new__(cls, name, bases, attrs):
+            _member_names_ = cls._member_names_
-        _choices = []
+            _member_map_ = cls._member_map_
-        collected = set()
+            _value2label_map_ = cls._value2label_map_
-        new_attrs = {}
-        for k, v in attrs.items():
-            if isinstance(v, tuple):
-                v = Choice(*v)
-                assert v not in collected, 'Cannot be defined repeatedly'
-                _choices.append(v)
-                collected.add(v)
-            new_attrs[k] = v
-        for base in bases:
-            if hasattr(base, '_choices'):
-                for c in base._choices:
-                    if c not in collected:
-                        _choices.append(c)
-                        collected.add(c)
-        new_attrs['_choices'] = _choices
-        new_attrs['_choices_dict'] = {c: c.label for c in _choices}
-        return type.__new__(cls, name, bases, new_attrs)
 
-    def __contains__(self, item):
+            for name in _member_names_:
-        return self._choices_dict.__contains__(item)
+                value = str(_member_map_[name])
-
+                label = _value2label_map_[value]
-    def __getitem__(self, item):
+                attrs[name] = value, label
-        return self._choices_dict.__getitem__(item)
+        bases = (models.TextChoices,)
-
+        return type(classname, bases, attrs)
-    def get(self, item, default=None):
-        return self._choices_dict.get(item, default)
-
-    @property
-    def choices(self):
-        return [(c, c.label) for c in self._choices]
-
-
-class ChoiceSet(metaclass=ChoiceSetType):
-    choices = None  # 用于 Django Model 中的 choices 配置， 为了代码提示在此声明
 
 
 class BitOperationChoice:
@@ -107,18 +84,18 @@ class BitOperationChoice:
         return [(cls.NAME_MAP[i], j) for i, j in cls.DB_CHOICES]
 
 
-class JMSBaseModel(Model):
+class BaseCreateUpdateModel(models.Model):
-    created_by = CharField(max_length=32, null=True, blank=True, verbose_name=_('Created by'))
+    created_by = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Created by'))
-    updated_by = CharField(max_length=32, null=True, blank=True, verbose_name=_('Updated by'))
+    updated_by = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Updated by'))
-    date_created = DateTimeField(auto_now_add=True, null=True, blank=True, verbose_name=_('Date created'))
+    date_created = models.DateTimeField(auto_now_add=True, null=True, blank=True, verbose_name=_('Date created'))
-    date_updated = DateTimeField(auto_now=True, verbose_name=_('Date updated'))
+    date_updated = models.DateTimeField(auto_now=True, verbose_name=_('Date updated'))
 
     class Meta:
         abstract = True
 
 
-class JMSModel(JMSBaseModel):
+class JMSBaseModel(BaseCreateUpdateModel):
-    id = UUIDField(default=uuid.uuid4, primary_key=True)
+    id = models.UUIDField(default=uuid.uuid4, primary_key=True)
 
     class Meta:
         abstract = True
@@ -129,7 +106,7 @@ def concated_display(name1, name2):
 
 
 def output_as_string(field_name):
-    return ExpressionWrapper(F(field_name), output_field=CharField())
+    return ExpressionWrapper(F(field_name), output_field=models.CharField())
 
 
 class UnionQuerySet(QuerySet):

apps/notifications/models/notification.py

@@ -1,11 +1,11 @@
 from django.db import models
 
-from common.db.models import JMSModel
+from common.db.models import JMSBaseModel
 
 __all__ = ('SystemMsgSubscription', 'UserMsgSubscription')
 
 
-class UserMsgSubscription(JMSModel):
+class UserMsgSubscription(JMSBaseModel):
     user = models.OneToOneField('users.User', related_name='user_msg_subscription', on_delete=models.CASCADE)
     receive_backends = models.JSONField(default=list)
 
@@ -13,7 +13,7 @@ class UserMsgSubscription(JMSModel):
         return f'{self.user} subscription: {self.receive_backends}'
 
 
-class SystemMsgSubscription(JMSModel):
+class SystemMsgSubscription(JMSBaseModel):
     message_type = models.CharField(max_length=128, unique=True)
     users = models.ManyToManyField('users.User', related_name='system_msg_subscriptions')
     groups = models.ManyToManyField('users.UserGroup', related_name='system_msg_subscriptions')

apps/notifications/models/site_msg.py

@@ -1,18 +1,18 @@
 from django.db import models
 
-from common.db.models import JMSModel
+from common.db.models import JMSBaseModel
 
 __all__ = ('SiteMessageUsers', 'SiteMessage')
 
 
-class SiteMessageUsers(JMSModel):
+class SiteMessageUsers(JMSBaseModel):
     sitemessage = models.ForeignKey('notifications.SiteMessage', on_delete=models.CASCADE, db_constraint=False, related_name='m2m_sitemessageusers')
     user = models.ForeignKey('users.User', on_delete=models.CASCADE, db_constraint=False, related_name='m2m_sitemessageusers')
     has_read = models.BooleanField(default=False)
     read_at = models.DateTimeField(default=None, null=True)
 
 
-class SiteMessage(JMSModel):
+class SiteMessage(JMSBaseModel):
     subject = models.CharField(max_length=1024)
     message = models.TextField()
     users = models.ManyToManyField(

apps/perms/models/asset_permission.py

@@ -2,10 +2,11 @@ import logging
 
 from django.utils.translation import ugettext_lazy as _
 from django.db.models import F, TextChoices
+from django.db import models
 
 from orgs.mixins.models import OrgModelMixin
-from common.db import models
 from common.utils import lazyproperty
+from common.db.models import BaseCreateUpdateModel
 from assets.models import Asset, SystemUser, Node, FamilyMixin
 
 from .base import BasePermission
@@ -89,7 +90,7 @@ class AssetPermission(BasePermission):
         return names
 
 
-class UserAssetGrantedTreeNodeRelation(OrgModelMixin, FamilyMixin, models.JMSBaseModel):
+class UserAssetGrantedTreeNodeRelation(OrgModelMixin, FamilyMixin, BaseCreateUpdateModel):
     class NodeFrom(TextChoices):
         granted = 'granted', 'Direct node granted'
         child = 'child', 'Have children node'

apps/rbac/models/role.py

@@ -1,7 +1,7 @@
 from django.utils.translation import ugettext_lazy as _, gettext
 from django.db import models
 
-from common.db.models import JMSModel
+from common.db.models import JMSBaseModel
 from common.utils import lazyproperty
 from .permission import Permission
 from ..builtin import BuiltinRole
@@ -22,7 +22,7 @@ class OrgRoleManager(models.Manager):
         return queryset.filter(scope=const.Scope.org)
 
 
-class Role(JMSModel):
+class Role(JMSBaseModel):
     """ 定义 角色 ｜ 角色-权限 关系 """
     Scope = const.Scope
 

apps/rbac/models/rolebinding.py

@@ -4,7 +4,7 @@ from django.db.models import Q
 from django.core.exceptions import ValidationError
 from rest_framework.serializers import ValidationError
 
-from common.db.models import JMSModel
+from common.db.models import JMSBaseModel
 from common.utils import lazyproperty
 from orgs.utils import current_org
 from .role import Role
@@ -29,7 +29,7 @@ class RoleBindingManager(models.Manager):
         return self.get_queryset()
 
 
-class RoleBinding(JMSModel):
+class RoleBinding(JMSBaseModel):
     Scope = Scope
     """ 定义 用户-角色 关系 """
     scope = models.CharField(

apps/terminal/migrations/0048_auto_20220407_1726.py

@@ -0,0 +1,18 @@
+# Generated by Django 3.1.14 on 2022-04-07 09:26
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('terminal', '0047_auto_20220302_1951'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='session',
+            name='protocol',
+            field=models.CharField(choices=[('ssh', 'SSH'), ('rdp', 'RDP'), ('telnet', 'Telnet'), ('vnc', 'VNC'), ('mysql', 'MySQL'), ('mariadb', 'MariaDB'), ('oracle', 'Oracle'), ('postgresql', 'PostgreSQL'), ('sqlserver', 'SQLServer'), ('redis', 'Redis'), ('mongodb', 'MongoDB'), ('k8s', 'K8S')], db_index=True, default='ssh', max_length=16),
+        ),
+    ]

apps/terminal/models/session.py

@@ -11,6 +11,7 @@ from django.core.files.storage import default_storage
 from django.core.cache import cache
 
 from assets.models import Asset
+from assets.const import Protocol
 from users.models import User
 from orgs.mixins.models import OrgModelMixin
 from django.db.models import TextChoices
@@ -24,20 +25,6 @@ class Session(OrgModelMixin):
         WT = 'WT', 'Web Terminal'
         DT = 'DT', 'DB Terminal'
 
-    class PROTOCOL(TextChoices):
-        SSH = 'ssh', 'ssh'
-        RDP = 'rdp', 'rdp'
-        VNC = 'vnc', 'vnc'
-        TELNET = 'telnet', 'telnet'
-        MYSQL = 'mysql', 'mysql'
-        ORACLE = 'oracle', 'oracle'
-        MARIADB = 'mariadb', 'mariadb'
-        SQLSERVER = 'sqlserver', 'sqlserver'
-        POSTGRESQL = 'postgresql', 'postgresql'
-        REDIS = 'redis', 'redis'
-        MONGODB = 'mongodb', 'MongoDB'
-        K8S = 'k8s', 'kubernetes'
-
     id = models.UUIDField(default=uuid.uuid4, primary_key=True)
     user = models.CharField(max_length=128, verbose_name=_("User"), db_index=True)
     user_id = models.CharField(blank=True, default='', max_length=36, db_index=True)
@@ -52,7 +39,7 @@ class Session(OrgModelMixin):
     has_replay = models.BooleanField(default=False, verbose_name=_("Replay"))
     has_command = models.BooleanField(default=False, verbose_name=_("Command"))
     terminal = models.ForeignKey('terminal.Terminal', null=True, on_delete=models.DO_NOTHING, db_constraint=False)
-    protocol = models.CharField(choices=PROTOCOL.choices, default='ssh', max_length=16, db_index=True)
+    protocol = models.CharField(choices=Protocol.choices, default='ssh', max_length=16, db_index=True)
     date_start = models.DateTimeField(verbose_name=_("Date start"), db_index=True, default=timezone.now)
     date_end = models.DateTimeField(verbose_name=_("Date end"), null=True)
 
@@ -131,29 +118,20 @@ class Session(OrgModelMixin):
 
     @property
     def can_join(self):
-        _PROTOCOL = self.PROTOCOL
         if self.is_finished:
             return False
         if self.login_from == self.LOGIN_FROM.RT:
             return False
-        if self.protocol in [
+        if Protocol in [
-            _PROTOCOL.SSH, _PROTOCOL.VNC, _PROTOCOL.RDP,
+            Protocol.SSH, Protocol.VNC, Protocol.RDP,
-            _PROTOCOL.TELNET, _PROTOCOL.K8S
+            Protocol.TELNET, Protocol.K8S
         ]:
             return True
         else:
             return False
 
-    @property
-    def db_protocols(self):
-        _PROTOCOL = self.PROTOCOL
-        return [_PROTOCOL.MYSQL, _PROTOCOL.MARIADB, _PROTOCOL.ORACLE,
-                _PROTOCOL.POSTGRESQL, _PROTOCOL.SQLSERVER,
-                _PROTOCOL.REDIS, _PROTOCOL.MONGODB]
-
     @property
     def can_terminate(self):
-        _PROTOCOL = self.PROTOCOL
         if self.is_finished:
             return False
         else:

utils/generate_fake_data/resources/assets.py

@@ -5,6 +5,7 @@ import forgery_py
 from .base import FakeDataGenerator
 
 from assets.models import *
+from assets.const import Protocol
 
 
 class AdminUsersGenerator(FakeDataGenerator):
@@ -28,7 +29,7 @@ class AdminUsersGenerator(FakeDataGenerator):
 class SystemUsersGenerator(FakeDataGenerator):
     def do_generate(self, batch, batch_size):
         system_users = []
-        protocols = list(dict(SystemUser.Protocol.choices).keys())
+        protocols = list(dict(Protocol.choices).keys())
         for i in batch:
             username = forgery_py.internet.user_name(True)
             protocol = random.choice(protocols)