jumpserver/apps/assets/task_handlers/backup/handlers.py

import os
import time
import pandas as pd
from collections import defaultdict

from django.conf import settings
from django.utils.translation import ugettext_lazy as _

from assets.models import AuthBook, Asset, BaseUser, ProtocolsMixin
from assets.notifications import AccountBackupExecutionTaskMsg
from applications.models import Account, Application
from applications.const import AppType
from users.models import User
from common.utils import get_logger
from common.utils.timezone import local_now_display
from common.utils.file import encrypt_and_compress_zip_file

logger = get_logger(__file__)

PATH = os.path.join(os.path.dirname(settings.BASE_DIR), 'tmp')


class AssetAccountHandler:
    @staticmethod
    def get_filename(plan_name):
        filename = os.path.join(
            PATH, f'{plan_name}-{_("Asset")}-{local_now_display()}-{time.time()}.xlsx'
        )
        return filename

    @staticmethod
    def create_df():
        df_dict = defaultdict(list)
        label_key = AuthBook._meta.verbose_name
        accounts = AuthBook.objects.all().prefetch_related('systemuser', 'asset')
        for account in accounts:
            account.load_auth()
            protocol = account.asset.protocol
            protocol_label = getattr(ProtocolsMixin.Protocol, protocol).label
            row = {
                getattr(Asset, 'hostname').field.verbose_name: account.asset.hostname,
                getattr(Asset, 'ip').field.verbose_name: account.asset.ip,
            }
            secret_row = AccountBackupHandler.create_secret_row(account)
            row.update(secret_row)
            row.update({
                getattr(Asset, 'protocol').field.verbose_name: protocol_label,
                getattr(AuthBook, 'version').field.verbose_name: account.version
            })
            df_dict[label_key].append(row)
        for k, v in df_dict.items():
            df_dict[k] = pd.DataFrame(v)
        return df_dict


class AppAccountHandler:
    @staticmethod
    def get_filename(plan_name):
        filename = os.path.join(
            PATH, f'{plan_name}-{_("Application")}-{local_now_display()}-{time.time()}.xlsx'
        )
        return filename

    @staticmethod
    def create_df():
        df_dict = defaultdict(list)
        accounts = Account.objects.all().prefetch_related('systemuser', 'app')
        for account in accounts:
            account.load_auth()
            app_type = account.app.type
            if app_type == 'postgresql':
                label_key = getattr(AppType, 'pgsql').label
            else:
                label_key = getattr(AppType, app_type).label
            row = {
                getattr(Application, 'name').field.verbose_name: account.app.name,
                getattr(Application, 'attrs').field.verbose_name: account.app.attrs
            }
            secret_row = AccountBackupHandler.create_secret_row(account)
            row.update(secret_row)
            row.update({
                getattr(Account, 'version').field.verbose_name: account.version
            })
            df_dict[label_key].append(row)
        for k, v in df_dict.items():
            df_dict[k] = pd.DataFrame(v)
        return df_dict


HANDLER_MAP = {
    'asset': AssetAccountHandler,
    'application': AppAccountHandler
}


class AccountBackupHandler:
    def __init__(self, execution):
        self.execution = execution
        self.plan_name = self.execution.plan.name
        self.is_frozen = False  # 任务状态冻结标志

    def create_excel(self):
        logger.info(
            '\n'
            '\033[32m>>> 正在生成资产及应用相关备份信息文件\033[0m'
            ''
        )
        # Print task start date
        time_start = time.time()
        info = {}
        for account_type in self.execution.types:
            if account_type in HANDLER_MAP:
                account_handler = HANDLER_MAP[account_type]
                df = account_handler.create_df()
                filename = account_handler.get_filename(self.plan_name)
                info[filename] = df
        for filename, df_dict in info.items():
            with pd.ExcelWriter(filename) as w:
                for sheet, df in df_dict.items():
                    sheet = sheet.replace(' ', '-')
                    getattr(df, 'to_excel')(w, sheet_name=sheet, index=False)
        timedelta = round((time.time() - time_start), 2)
        logger.info('步骤完成: 用时 {}s'.format(timedelta))
        return list(info.keys())

    def send_backup_mail(self, files):
        recipients = self.execution.plan_snapshot.get('recipients')
        if not recipients:
            return
        recipients = User.objects.filter(id__in=list(recipients))
        logger.info(
            '\n'
            '\033[32m>>> 发送备份邮件\033[0m'
            ''
        )
        plan_name = self.plan_name
        for user in recipients:
            if not user.secret_key:
                attachment_list = []
            else:
                password = user.secret_key.encode('utf8')
                attachment = os.path.join(PATH, f'{plan_name}-{local_now_display()}-{time.time()}.zip')
                encrypt_and_compress_zip_file(attachment, password, files)
                attachment_list = [attachment, ]
            AccountBackupExecutionTaskMsg(plan_name, user).publish(attachment_list)
            logger.info('邮件已发送至{}({})'.format(user, user.email))
        for file in files:
            os.remove(file)

    def step_perform_task_update(self, is_success, reason):
        self.execution.reason = reason[:1024]
        self.execution.is_success = is_success
        self.execution.save()
        logger.info('已完成对任务状态的更新')

    def step_finished(self, is_success):
        if is_success:
            logger.info('任务执行成功')
        else:
            logger.error('任务执行失败')

    def _run(self):
        is_success = False
        error = '-'
        try:
            files = self.create_excel()
            self.send_backup_mail(files)
        except Exception as e:
            self.is_frozen = True
            logger.error('任务执行被异常中断')
            logger.info('下面打印发生异常的 Traceback 信息 : ')
            logger.error(e, exc_info=True)
            error = str(e)
        else:
            is_success = True
        finally:
            reason = error
            self.step_perform_task_update(is_success, reason)
            self.step_finished(is_success)

    def run(self):
        logger.info('任务开始: {}'.format(local_now_display()))
        time_start = time.time()
        try:
            self._run()
        except Exception as e:
            logger.error('任务运行出现异常')
            logger.error('下面显示异常 Traceback 信息: ')
            logger.error(e, exc_info=True)
        finally:
            logger.info('\n任务结束: {}'.format(local_now_display()))
            timedelta = round((time.time() - time_start), 2)
            logger.info('用时: {}'.format(timedelta))

    @staticmethod
    def create_secret_row(instance):
        row = {
            getattr(BaseUser, 'username').field.verbose_name: instance.username,
            getattr(BaseUser, 'password').field.verbose_name: instance.password,
            getattr(BaseUser, 'private_key').field.verbose_name: instance.private_key,
            getattr(BaseUser, 'public_key').field.verbose_name: instance.public_key
        }
        return row