mirror of https://github.com/jumpserver/jumpserver
204 lines
7.3 KiB
Python
204 lines
7.3 KiB
Python
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
|