perf: 优化账号备份

pull/8733/head
feng626 2022-09-01 16:31:20 +08:00 committed by Jiangjie.Bai
parent f0dc519423
commit ef4cc5f646
3 changed files with 139 additions and 59 deletions

View File

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

View File

@ -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'),

View File

@ -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