From ef4cc5f646a8437f33d13e7f090a03aec1f8d6b8 Mon Sep 17 00:00:00 2001 From: feng626 <1304903146@qq.com> Date: Thu, 1 Sep 2022 16:31:20 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=E8=B4=A6=E5=8F=B7?= =?UTF-8?q?=E5=A4=87=E4=BB=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/applications/serializers/application.py | 54 +++++--- apps/assets/serializers/account.py | 20 ++- apps/assets/task_handlers/backup/handlers.py | 124 +++++++++++++------ 3 files changed, 139 insertions(+), 59 deletions(-) diff --git a/apps/applications/serializers/application.py b/apps/applications/serializers/application.py index 35a07e262..ac8f38f2a 100644 --- a/apps/applications/serializers/application.py +++ b/apps/applications/serializers/application.py @@ -16,7 +16,7 @@ from .. import const __all__ = [ 'AppSerializer', 'MiniAppSerializer', 'AppSerializerMixin', - 'AppAccountSerializer', 'AppAccountSecretSerializer' + 'AppAccountSerializer', 'AppAccountSecretSerializer', 'AppAccountBackUpSerializer' ] @@ -32,21 +32,23 @@ class AppSerializerMixin(serializers.Serializer): return instance def get_attrs_serializer(self): - default_serializer = serializers.Serializer(read_only=True) instance = self.app - if instance: - _type = instance.type - _category = instance.category - else: - _type = self.context['request'].query_params.get('type') - _category = self.context['request'].query_params.get('category') - if _type: - if isinstance(self, AppAccountSecretSerializer): - serializer_class = type_secret_serializer_classes_mapping.get(_type) + tp = getattr(self, 'tp', None) + default_serializer = serializers.Serializer(read_only=True) + if not tp: + if instance: + tp = instance.type + category = instance.category else: - serializer_class = type_serializer_classes_mapping.get(_type) - elif _category: - serializer_class = category_serializer_classes_mapping.get(_category) + tp = self.context['request'].query_params.get('type') + category = self.context['request'].query_params.get('category') + if tp: + if isinstance(self, AppAccountBackUpSerializer): + serializer_class = type_secret_serializer_classes_mapping.get(tp) + else: + serializer_class = type_serializer_classes_mapping.get(tp) + elif category: + serializer_class = category_serializer_classes_mapping.get(category) else: serializer_class = default_serializer @@ -154,11 +156,6 @@ class AppAccountSerializer(AppSerializerMixin, AuthSerializerMixin, BulkOrgResou class AppAccountSecretSerializer(SecretReadableMixin, AppAccountSerializer): class Meta(AppAccountSerializer.Meta): - fields_backup = [ - 'id', 'app_display', 'attrs', 'username', 'password', 'private_key', - 'public_key', 'date_created', 'date_updated', 'version' - ] - extra_kwargs = { 'password': {'write_only': False}, 'private_key': {'write_only': False}, @@ -166,3 +163,22 @@ class AppAccountSecretSerializer(SecretReadableMixin, AppAccountSerializer): 'app_display': {'label': _('Application display')}, 'systemuser_display': {'label': _('System User')} } + + +class AppAccountBackUpSerializer(AppAccountSecretSerializer): + class Meta(AppAccountSecretSerializer.Meta): + fields = [ + 'id', 'app_display', 'attrs', 'username', 'password', 'private_key', + 'public_key', 'date_created', 'date_updated', 'version' + ] + + def __init__(self, *args, **kwargs): + self.tp = kwargs.pop('tp', None) + super().__init__(*args, **kwargs) + + @classmethod + def setup_eager_loading(cls, queryset): + return queryset + + def to_representation(self, instance): + return super(AppAccountSerializer, self).to_representation(instance) diff --git a/apps/assets/serializers/account.py b/apps/assets/serializers/account.py index e4e320f33..9b06944c3 100644 --- a/apps/assets/serializers/account.py +++ b/apps/assets/serializers/account.py @@ -76,10 +76,6 @@ class AccountSerializer(AuthSerializerMixin, BulkOrgResourceModelSerializer): class AccountSecretSerializer(SecretReadableMixin, AccountSerializer): class Meta(AccountSerializer.Meta): - fields_backup = [ - 'hostname', 'ip', 'platform', 'protocols', 'username', 'password', - 'private_key', 'public_key', 'date_created', 'date_updated', 'version' - ] extra_kwargs = { 'password': {'write_only': False}, 'private_key': {'write_only': False}, @@ -88,6 +84,22 @@ class AccountSecretSerializer(SecretReadableMixin, AccountSerializer): } +class AccountBackUpSerializer(AccountSecretSerializer): + class Meta(AccountSecretSerializer.Meta): + fields = [ + 'id', 'hostname', 'ip', 'username', 'password', + 'private_key', 'public_key', 'date_created', + 'date_updated', 'version' + ] + + @classmethod + def setup_eager_loading(cls, queryset): + return queryset + + def to_representation(self, instance): + return super(AccountSerializer, self).to_representation(instance) + + class AccountTaskSerializer(serializers.Serializer): ACTION_CHOICES = ( ('test', 'test'), diff --git a/apps/assets/task_handlers/backup/handlers.py b/apps/assets/task_handlers/backup/handlers.py index 311d8e395..a29058fac 100644 --- a/apps/assets/task_handlers/backup/handlers.py +++ b/apps/assets/task_handlers/backup/handlers.py @@ -4,15 +4,16 @@ from openpyxl import Workbook from collections import defaultdict, OrderedDict from django.conf import settings +from django.db.models import F from django.utils.translation import ugettext_lazy as _ from rest_framework import serializers -from assets.models import AuthBook -from assets.serializers import AccountSecretSerializer +from assets.models import AuthBook, SystemUser, Asset +from assets.serializers import AccountBackUpSerializer from assets.notifications import AccountBackupExecutionTaskMsg -from applications.models import Account +from applications.models import Account, Application from applications.const import AppType -from applications.serializers import AppAccountSecretSerializer +from applications.serializers import AppAccountBackUpSerializer from users.models import User from common.utils import get_logger from common.utils.timezone import local_now_display @@ -38,7 +39,7 @@ class BaseAccountHandler: @classmethod def get_header_fields(cls, serializer: serializers.Serializer): try: - backup_fields = getattr(serializer, 'Meta').fields_backup + backup_fields = getattr(serializer, 'Meta').fields except AttributeError: backup_fields = serializer.fields.keys() header_fields = {} @@ -51,17 +52,41 @@ class BaseAccountHandler: header_fields[field] = str(v.label) return header_fields + @staticmethod + def load_auth(tp, value, system_user): + if value: + return value + if system_user: + return getattr(system_user, tp, '') + return '' + @classmethod - def create_row(cls, account, serializer_cls, header_fields=None): - serializer = serializer_cls(account) - if not header_fields: - header_fields = cls.get_header_fields(serializer) - data = cls.unpack_data(serializer.data) + def replace_auth(cls, account, system_user_dict): + system_user = system_user_dict.get(account.systemuser_id) + account.username = cls.load_auth('username', account.username, system_user) + account.password = cls.load_auth('password', account.password, system_user) + account.private_key = cls.load_auth('private_key', account.private_key, system_user) + account.public_key = cls.load_auth('public_key', account.public_key, system_user) + return account + + @classmethod + def create_row(cls, data, header_fields): + data = cls.unpack_data(data) row_dict = {} for field, header_name in header_fields.items(): - row_dict[header_name] = str(data[field]) + row_dict[header_name] = str(data.get(field, field)) return row_dict + @classmethod + def add_rows(cls, data, header_fields, sheet): + data_map = defaultdict(list) + for i in data: + row = cls.create_row(i, header_fields) + if sheet not in data_map: + data_map[sheet].append(list(row.keys())) + data_map[sheet].append(list(row.values())) + return data_map + class AssetAccountHandler(BaseAccountHandler): @staticmethod @@ -72,22 +97,27 @@ class AssetAccountHandler(BaseAccountHandler): return filename @classmethod - def create_data_map(cls): - data_map = defaultdict(list) + def replace_account_info(cls, account, asset_dict, system_user_dict): + asset = asset_dict.get(account.asset_id) + account.ip = asset.ip if asset else '' + account.hostname = asset.hostname if asset else '' + account = cls.replace_auth(account, system_user_dict) + return account + + @classmethod + def create_data_map(cls, system_user_dict): sheet_name = AuthBook._meta.verbose_name + assets = Asset.objects.only('id', 'hostname', 'ip') + asset_dict = {asset.id: asset for asset in assets} + accounts = AuthBook.objects.all() + if not accounts.exists(): + return - accounts = AuthBook.get_queryset().select_related('systemuser') - if not accounts.first(): - return data_map - - header_fields = cls.get_header_fields(AccountSecretSerializer(accounts.first())) + header_fields = cls.get_header_fields(AccountBackUpSerializer(accounts.first())) for account in accounts: - account.load_auth() - row = cls.create_row(account, AccountSecretSerializer, header_fields) - if sheet_name not in data_map: - data_map[sheet_name].append(list(row.keys())) - data_map[sheet_name].append(list(row.values())) - + cls.replace_account_info(account, asset_dict, system_user_dict) + data = AccountBackUpSerializer(accounts, many=True).data + data_map = cls.add_rows(data, header_fields, sheet_name) logger.info('\n\033[33m- 共收集 {} 条资产账号\033[0m'.format(accounts.count())) return data_map @@ -101,18 +131,36 @@ class AppAccountHandler(BaseAccountHandler): return filename @classmethod - def create_data_map(cls): - data_map = defaultdict(list) - accounts = Account.get_queryset().select_related('systemuser') - for account in accounts: - account.load_auth() - app_type = account.type + def replace_account_info(cls, account, app_dict, system_user_dict): + app = app_dict.get(account.app_id) + account.type = app.type if app else '' + account.app_display = app.name if app else '' + account.category = app.category if app else '' + account = cls.replace_auth(account, system_user_dict) + return account + + @classmethod + def create_data_map(cls, system_user_dict): + apps = Application.objects.only('id', 'type', 'name', 'category') + app_dict = {app.id: app for app in apps} + qs = Account.objects.all().annotate(app_type=F('app__type')) + if not qs.exists(): + return + + account_type_map = defaultdict(list) + for i in qs: + account_type_map[i.app_type].append(i) + data_map = {} + for app_type, accounts in account_type_map.items(): sheet_name = AppType.get_label(app_type) - row = cls.create_row(account, AppAccountSecretSerializer) - if sheet_name not in data_map: - data_map[sheet_name].append(list(row.keys())) - data_map[sheet_name].append(list(row.values())) - logger.info('\n\033[33m- 共收集{}条应用账号\033[0m'.format(accounts.count())) + header_fields = cls.get_header_fields(AppAccountBackUpSerializer(tp=app_type)) + if not accounts: + continue + for account in accounts: + cls.replace_account_info(account, app_dict, system_user_dict) + data = AppAccountBackUpSerializer(accounts, many=True, app_type=app_type).data + data_map.update(cls.add_rows(data, header_fields, sheet_name)) + logger.info('\n\033[33m- 共收集{}条应用账号\033[0m'.format(qs.count())) return data_map @@ -137,12 +185,16 @@ class AccountBackupHandler: # Print task start date time_start = time.time() files = [] + system_user_qs = SystemUser.objects.only( + 'id', 'username', 'password', 'private_key', 'public_key' + ) + system_user_dict = {i.id: i for i in system_user_qs} for account_type in self.execution.types: handler = handler_map.get(account_type) if not handler: continue - data_map = handler.create_data_map() + data_map = handler.create_data_map(system_user_dict) if not data_map: continue