账号备份

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),
+    ]

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()

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)

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