From d418c28e988a418613f6e95774fe491e42ffc9b5 Mon Sep 17 00:00:00 2001 From: ibuler Date: Thu, 7 Apr 2022 18:51:35 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BF=AE=E6=94=B9=20base?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/acls/serializers/login_asset_acl.py | 8 +- apps/assets/api/system_user.py | 3 +- apps/assets/const.py | 94 ++++++++++++++++++- .../migrations/0097_auto_20220407_1726.py | 34 +++++++ .../migrations/0098_auto_20220407_1730.py | 23 +++++ apps/assets/models/asset/_category.py | 73 -------------- apps/assets/models/asset/common.py | 13 +-- apps/assets/models/asset/database.py | 1 - apps/assets/models/asset/host.py | 5 +- apps/assets/models/asset/network.py | 5 + apps/assets/models/platform.py | 12 +-- apps/assets/models/user.py | 18 +--- apps/assets/serializers/asset/common.py | 4 +- apps/assets/serializers/system_user.py | 17 ++-- apps/authentication/models.py | 7 +- apps/common/db/models.py | 79 ++++++---------- apps/notifications/models/notification.py | 6 +- apps/notifications/models/site_msg.py | 6 +- apps/perms/models/asset_permission.py | 5 +- apps/rbac/models/role.py | 4 +- apps/rbac/models/rolebinding.py | 4 +- .../migrations/0048_auto_20220407_1726.py | 18 ++++ apps/terminal/models/session.py | 32 +------ utils/generate_fake_data/resources/assets.py | 3 +- 24 files changed, 256 insertions(+), 218 deletions(-) create mode 100644 apps/assets/migrations/0097_auto_20220407_1726.py create mode 100644 apps/assets/migrations/0098_auto_20220407_1730.py delete mode 100644 apps/assets/models/asset/_category.py create mode 100644 apps/terminal/migrations/0048_auto_20220407_1726.py diff --git a/apps/acls/serializers/login_asset_acl.py b/apps/acls/serializers/login_asset_acl.py index df62a04f8..34d6711c8 100644 --- a/apps/acls/serializers/login_asset_acl.py +++ b/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]) ) ) diff --git a/apps/assets/api/system_user.py b/apps/assets/api/system_user.py index f679b4e3f..2739044de 100644 --- a/apps/assets/api/system_user.py +++ b/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) diff --git a/apps/assets/const.py b/apps/assets/const.py index ec51c5a2b..ccd48205b 100644 --- a/apps/assets/const.py +++ b/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, + ] + diff --git a/apps/assets/migrations/0097_auto_20220407_1726.py b/apps/assets/migrations/0097_auto_20220407_1726.py new file mode 100644 index 000000000..5459d8cee --- /dev/null +++ b/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'), + ), + ] diff --git a/apps/assets/migrations/0098_auto_20220407_1730.py b/apps/assets/migrations/0098_auto_20220407_1730.py new file mode 100644 index 000000000..a50986cbc --- /dev/null +++ b/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'), + ), + ] diff --git a/apps/assets/models/asset/_category.py b/apps/assets/models/asset/_category.py deleted file mode 100644 index 8b7f0fa0b..000000000 --- a/apps/assets/models/asset/_category.py +++ /dev/null @@ -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 - - - diff --git a/apps/assets/models/asset/common.py b/apps/assets/models/asset/common.py index c8fc0ce23..05404ff94 100644 --- a/apps/assets/models/asset/common.py +++ b/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 diff --git a/apps/assets/models/asset/database.py b/apps/assets/models/asset/database.py index ff0be89a7..84d174277 100644 --- a/apps/assets/models/asset/database.py +++ b/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) - diff --git a/apps/assets/models/asset/host.py b/apps/assets/models/asset/host.py index ddaa1766d..a729c2da6 100644 --- a/apps/assets/models/asset/host.py +++ b/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): diff --git a/apps/assets/models/asset/network.py b/apps/assets/models/asset/network.py index e69de29bb..42d35e06b 100644 --- a/apps/assets/models/asset/network.py +++ b/apps/assets/models/asset/network.py @@ -0,0 +1,5 @@ +from .common import Asset + + +class Network(Asset): + pass diff --git a/apps/assets/models/platform.py b/apps/assets/models/platform.py index 5b8aafe7e..63ea42ff1 100644 --- a/apps/assets/models/platform.py +++ b/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")) diff --git a/apps/assets/models/user.py b/apps/assets/models/user.py index ce0029768..e179c9b22 100644 --- a/apps/assets/models/user.py +++ b/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 - - 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' + Protocol = Protocol 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')) diff --git a/apps/assets/serializers/asset/common.py b/apps/assets/serializers/asset/common.py index 94283e1eb..94e5be260 100644 --- a/apps/assets/serializers/asset/common.py +++ b/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 + [ diff --git a/apps/assets/serializers/system_user.py b/apps/assets/serializers/system_user.py index 850870762..4ede6ac65 100644 --- a/apps/assets/serializers/system_user.py +++ b/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 \ - and protocol != SystemUser.Protocol.redis: + if login_mode == SystemUser.LOGIN_AUTO and protocol != Protocol.vnc \ + 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) - if protocol not in [SystemUser.Protocol.ssh.value]: + protocol = self.get_initial_value('protocol', default=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 diff --git a/apps/authentication/models.py b/apps/authentication/models.py index 1b353f737..3b469eb95 100644 --- a/apps/authentication/models.py +++ b/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 可能要授权给 普通用户, 或者放开就行 diff --git a/apps/common/db/models.py b/apps/common/db/models.py index 2989f734e..a871de31c 100644 --- a/apps/common/db/models.py +++ b/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): - def __new__(cls, value, label=''): # `deepcopy` 的时候不会传 `label` - self = super().__new__(cls, value) - self.label = label - return self +class IncludesTextChoicesMeta(type): + def __new__(metacls, classname, bases, classdict): + includes = classdict.pop('includes', None) + assert includes + attrs = _EnumDict() + for k, v in classdict.items(): + attrs[k] = v -class ChoiceSetType(type): - def __new__(cls, name, bases, attrs): - _choices = [] - collected = set() - 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) + for cls in includes: + _member_names_ = cls._member_names_ + _member_map_ = cls._member_map_ + _value2label_map_ = cls._value2label_map_ - def __contains__(self, item): - return self._choices_dict.__contains__(item) - - def __getitem__(self, item): - return self._choices_dict.__getitem__(item) - - 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 配置， 为了代码提示在此声明 + for name in _member_names_: + value = str(_member_map_[name]) + label = _value2label_map_[value] + attrs[name] = value, label + bases = (models.TextChoices,) + return type(classname, bases, attrs) class BitOperationChoice: @@ -107,18 +84,18 @@ class BitOperationChoice: return [(cls.NAME_MAP[i], j) for i, j in cls.DB_CHOICES] -class JMSBaseModel(Model): - created_by = 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')) - date_created = DateTimeField(auto_now_add=True, null=True, blank=True, verbose_name=_('Date created')) - date_updated = DateTimeField(auto_now=True, verbose_name=_('Date updated')) +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')) + date_created = models.DateTimeField(auto_now_add=True, null=True, blank=True, verbose_name=_('Date created')) + date_updated = models.DateTimeField(auto_now=True, verbose_name=_('Date updated')) class Meta: abstract = True -class JMSModel(JMSBaseModel): - id = UUIDField(default=uuid.uuid4, primary_key=True) +class JMSBaseModel(BaseCreateUpdateModel): + 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): diff --git a/apps/notifications/models/notification.py b/apps/notifications/models/notification.py index d50168576..a3095fbe5 100644 --- a/apps/notifications/models/notification.py +++ b/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') diff --git a/apps/notifications/models/site_msg.py b/apps/notifications/models/site_msg.py index 3e3c09baa..e08cd5c71 100644 --- a/apps/notifications/models/site_msg.py +++ b/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( diff --git a/apps/perms/models/asset_permission.py b/apps/perms/models/asset_permission.py index ea795d889..b258e6742 100644 --- a/apps/perms/models/asset_permission.py +++ b/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' diff --git a/apps/rbac/models/role.py b/apps/rbac/models/role.py index b47bccb5b..631d979b7 100644 --- a/apps/rbac/models/role.py +++ b/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 diff --git a/apps/rbac/models/rolebinding.py b/apps/rbac/models/rolebinding.py index a2ee06022..a871b4afe 100644 --- a/apps/rbac/models/rolebinding.py +++ b/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( diff --git a/apps/terminal/migrations/0048_auto_20220407_1726.py b/apps/terminal/migrations/0048_auto_20220407_1726.py new file mode 100644 index 000000000..1eb432efc --- /dev/null +++ b/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), + ), + ] diff --git a/apps/terminal/models/session.py b/apps/terminal/models/session.py index 09f1c9e91..e1a65ceaa 100644 --- a/apps/terminal/models/session.py +++ b/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 [ - _PROTOCOL.SSH, _PROTOCOL.VNC, _PROTOCOL.RDP, - _PROTOCOL.TELNET, _PROTOCOL.K8S + if Protocol in [ + Protocol.SSH, Protocol.VNC, Protocol.RDP, + 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: diff --git a/utils/generate_fake_data/resources/assets.py b/utils/generate_fake_data/resources/assets.py index 3aa11cc99..de97a27a9 100644 --- a/utils/generate_fake_data/resources/assets.py +++ b/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)