diff --git a/apps/assets/migrations/0107_alter_accountbackupplan_types.py b/apps/assets/migrations/0107_alter_accountbackupplan_types.py new file mode 100644 index 000000000..2ce2d9302 --- /dev/null +++ b/apps/assets/migrations/0107_alter_accountbackupplan_types.py @@ -0,0 +1,59 @@ +# Generated by Django 3.2.13 on 2022-08-29 11:46 +from django.db import migrations, models + +from assets.const import Category +from assets.models import Type + + +def update_account_backup_type(apps, schema_editor): + backup_model = apps.get_model('assets', 'AccountBackupPlan') + all_number = 4294967295 + asset_number = Type.choices_to_value([Category.HOST]) + app_number = Type.choices_to_value([ + Category.NETWORKING, Category.DATABASE, Category.CLOUD, Category.WEB] + ) + + backup_model.objects.filter(types=255).update(types=all_number) + backup_model.objects.filter(types=1).update(types=asset_number) + backup_model.objects.filter(types=2).update(types=app_number) + + backup_execution_model = apps.get_model('assets', 'AccountBackupPlanExecution') + choices_dict = { + 'all': Type.get_types(value=all_number), + 'asset': Type.get_types(value=asset_number), + 'app': Type.get_types(value=app_number) + } + qs_dict = { + 'all': backup_execution_model.objects.filter(plan__types=255), + 'asset': backup_execution_model.objects.filter(plan__types=1), + 'app': backup_execution_model.objects.filter(plan__types=2) + } + + backup_executions = [] + for k, qs in qs_dict.items(): + type_choices = choices_dict[k] + for i in qs: + i.plan_snapshot['types'] = type_choices + backup_executions.append(i) + backup_execution_model.objects.bulk_update(backup_executions, fields=['plan_snapshot']) + + +class Migration(migrations.Migration): + dependencies = [ + ('assets', '0106_auto_20220819_1523'), + ] + + operations = [ + migrations.AlterField( + model_name='accountbackupplan', + name='types', + field=models.BigIntegerField( + choices=[(4294967295, 'All'), (1, 'Linux'), (2, 'Windows'), (4, 'Unix'), (8, 'BSD'), (16, 'MacOS'), + (32, 'Mainframe'), (64, 'Other host'), (127, 'Host'), (128, 'Switch'), (256, 'Router'), + (512, 'Firewall'), (1024, 'Other device'), (1920, 'NetworkDevice'), (2048, 'MySQL'), + (4096, 'MariaDB'), (8192, 'PostgreSQL'), (16384, 'Oracle'), (32768, 'SQLServer'), + (65536, 'MongoDB'), (131072, 'Redis'), (260096, 'Database'), (262144, 'Clouding'), + (524288, 'Web')], default=4294967295, verbose_name='Type'), + ), + migrations.RunPython(update_account_backup_type), + ] diff --git a/apps/assets/models/backup.py b/apps/assets/models/backup.py index 437e91dbd..e27564d9a 100644 --- a/apps/assets/models/backup.py +++ b/apps/assets/models/backup.py @@ -2,6 +2,7 @@ # -*- coding: utf-8 -*- # import uuid +from functools import reduce from celery import current_task from django.db import models @@ -13,40 +14,63 @@ from common.utils import get_logger from common.db.encoder import ModelJSONFieldEncoder from common.db.models import BitOperationChoice from common.mixins.models import CommonModelMixin +from ..const import AllTypes, Category __all__ = ['AccountBackupPlan', 'AccountBackupPlanExecution', 'Type'] logger = get_logger(__file__) +def _choice_map(default=None): + offset = 0 + temp_key = 0b1 + + if default is None: + _all = (0b1 << 32) - 1 + else: + _all = default + + choices = { + _all: ('all', 'All') + } + + for info in AllTypes.grouped_choices_to_objs(): + temp_keys = [] + for c in info['children']: + key = temp_key << offset + temp_keys.append(key) + choices[key] = (c['value'], c['display_name']) + offset += 1 + parent_key = reduce(lambda x, y: x | y, temp_keys) + choices[parent_key] = (info['value'], info['display_name']) + return choices + + class Type(BitOperationChoice): NONE = 0 - ALL = 0xff - Asset = 0b1 - App = 0b1 << 1 + ALL = (0b1 << 32) - 1 + TYPE_MAP = _choice_map(ALL) - DB_CHOICES = ( - (ALL, _('All')), - (Asset, _('Asset')), - (App, _('Application')) - ) + DB_CHOICES = tuple((k, v[1]) for k, v in TYPE_MAP.items()) - NAME_MAP = { - ALL: "all", - Asset: "asset", - App: "application" - } + NAME_MAP = {k: v[0] for k, v in TYPE_MAP.items()} NAME_MAP_REVERSE = {v: k for k, v in NAME_MAP.items()} CHOICES = [] for i, j in DB_CHOICES: CHOICES.append((NAME_MAP[i], j)) + @classmethod + def get_types(cls, value: int) -> list: + exclude_types = ['all'] + Category.values + current_all = cls.value_to_choices(value) + return list(filter(lambda x: x not in exclude_types, current_all)) + class AccountBackupPlan(CommonModelMixin, PeriodTaskModelMixin, OrgModelMixin): id = models.UUIDField(default=uuid.uuid4, primary_key=True) - types = models.IntegerField(choices=Type.DB_CHOICES, default=Type.ALL, verbose_name=_('Type')) + types = models.BigIntegerField(choices=Type.DB_CHOICES, default=Type.ALL, verbose_name=_('Type')) recipients = models.ManyToManyField( 'users.User', related_name='recipient_escape_route_plans', blank=True, verbose_name=_("Recipient") @@ -77,7 +101,7 @@ class AccountBackupPlan(CommonModelMixin, PeriodTaskModelMixin, OrgModelMixin): 'crontab': self.crontab, 'org_id': self.org_id, 'created_by': self.created_by, - 'types': Type.value_to_choices(self.types), + 'types': Type.get_types(self.types), 'recipients': { str(recipient.id): (str(recipient), bool(recipient.secret_key)) for recipient in self.recipients.all() diff --git a/apps/assets/serializers/account/account_template.py b/apps/assets/serializers/account/account_template.py index dbe298f43..068239e74 100644 --- a/apps/assets/serializers/account/account_template.py +++ b/apps/assets/serializers/account/account_template.py @@ -31,7 +31,6 @@ class AccountTemplateSerializer(AuthSerializerMixin, BulkOrgResourceModelSeriali for k, v in cls().fields.items(): if v.required and k not in attrs: required_field_dict[k] = error - print(required_field_dict) if not required_field_dict: return raise serializers.ValidationError(required_field_dict) diff --git a/apps/assets/task_handlers/backup/handlers.py b/apps/assets/task_handlers/backup/handlers.py index df7bb52da..9d0019054 100644 --- a/apps/assets/task_handlers/backup/handlers.py +++ b/apps/assets/task_handlers/backup/handlers.py @@ -4,10 +4,10 @@ from openpyxl import Workbook from collections import defaultdict, OrderedDict from django.conf import settings -from django.utils.translation import ugettext_lazy as _ +from django.db.models import F from rest_framework import serializers -from assets.models import Account +from assets.models import Account, Type from assets.serializers import AccountSecretSerializer from assets.notifications import AccountBackupExecutionTaskMsg from users.models import User @@ -64,34 +64,33 @@ class AssetAccountHandler(BaseAccountHandler): @staticmethod def get_filename(plan_name): filename = os.path.join( - PATH, f'{plan_name}-{_("Asset")}-{local_now_display()}-{time.time()}.xlsx' + PATH, f'{plan_name}-{local_now_display()}-{time.time()}.xlsx' ) return filename @classmethod - def create_data_map(cls): + def create_data_map(cls, types: list): data_map = defaultdict(list) - sheet_name = Account._meta.verbose_name - accounts = Account.objects.all() + # TODO 可以优化一下查询 在账号上做type的缓存 避免数据量大时连表操作 + accounts = Account.objects.filter( + asset__platform__in=types + ).annotate(type=F('asset__platform__type')) if not accounts.first(): return data_map + type_dict = dict(Type.CHOICES) header_fields = cls.get_header_fields(AccountSecretSerializer(accounts.first())) for account in accounts: - account.load_auth() + sheet_name = type_dict[account.type] 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())) - logger.info('\n\033[33m- 共收集 {} 条资产账号\033[0m'.format(accounts.count())) + logger.info('\n\033[33m- 共收集 {} 条账号\033[0m'.format(accounts.count())) return data_map -handler_map = { - 'asset': AssetAccountHandler, -} - class AccountBackupHandler: def __init__(self, execution): @@ -108,24 +107,21 @@ class AccountBackupHandler: # Print task start date time_start = time.time() files = [] - for account_type in self.execution.types: - handler = handler_map.get(account_type) - if not handler: - continue + types = self.execution.types - data_map = handler.create_data_map() - if not data_map: - continue + data_map = AssetAccountHandler.create_data_map(types) + if not data_map: + return files - filename = handler.get_filename(self.plan_name) + filename = AssetAccountHandler.get_filename(self.plan_name) - wb = Workbook(filename) - for sheet, data in data_map.items(): - ws = wb.create_sheet(str(sheet)) - for row in data: - ws.append(row) - wb.save(filename) - files.append(filename) + wb = Workbook(filename) + for sheet, data in data_map.items(): + ws = wb.create_sheet(str(sheet)) + for row in data: + ws.append(row) + wb.save(filename) + files.append(filename) timedelta = round((time.time() - time_start), 2) logger.info('步骤完成: 用时 {}s'.format(timedelta)) return files