import os import time from copy import deepcopy from openpyxl import Workbook from collections import defaultdict from django.conf import settings from django.utils import timezone from users.models import User from accounts.models import ChangeSecretRecord from accounts.notifications import ChangeSecretExecutionTaskMsg from accounts.serializers import ChangeSecretRecordBackUpSerializer from accounts.const import AutomationTypes, SecretType, SSHKeyStrategy, SecretStrategy from common.utils import get_logger, lazyproperty from common.utils.file import encrypt_and_compress_zip_file from common.utils.timezone import local_now_display from ...utils import SecretGenerator from ..base.manager import AccountBasePlaybookManager logger = get_logger(__name__) class ChangeSecretManager(AccountBasePlaybookManager): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.method_hosts_mapper = defaultdict(list) self.secret_type = self.execution.snapshot['secret_type'] self.secret_strategy = self.execution.snapshot.get( 'secret_strategy', SecretStrategy.custom ) self.ssh_key_change_strategy = self.execution.snapshot.get( 'ssh_key_change_strategy', SSHKeyStrategy.add ) self.snapshot_account_usernames = self.execution.snapshot['accounts'] self._password_generated = None self._ssh_key_generated = None self.name_recorder_mapper = {} # 做个映射,方便后面处理 @classmethod def method_type(cls): return AutomationTypes.change_secret @lazyproperty def related_accounts(self): pass def get_kwargs(self, account, secret): kwargs = {} if self.secret_type != SecretType.SSH_KEY: return kwargs kwargs['strategy'] = self.ssh_key_change_strategy kwargs['exclusive'] = 'yes' if kwargs['strategy'] == SSHKeyStrategy.set else 'no' if kwargs['strategy'] == SSHKeyStrategy.set_jms: kwargs['dest'] = '/home/{}/.ssh/authorized_keys'.format(account.username) kwargs['regexp'] = '.*{}$'.format(secret.split()[2].strip()) return kwargs @lazyproperty def secret_generator(self): return SecretGenerator( self.secret_strategy, self.secret_type, self.execution.snapshot.get('password_rules') ) def get_secret(self): if self.secret_strategy == SecretStrategy.custom: return self.execution.snapshot['secret'] else: return self.secret_generator.get_secret() def host_callback(self, host, asset=None, account=None, automation=None, path_dir=None, **kwargs): host = super().host_callback(host, asset=asset, account=account, automation=automation, **kwargs) if host.get('error'): return host accounts = asset.accounts.all() if account: accounts = accounts.exclude(username=account.username) if '*' not in self.snapshot_account_usernames: accounts = accounts.filter(username__in=self.snapshot_account_usernames) accounts = accounts.filter(secret_type=self.secret_type) method_attr = getattr(automation, self.method_type() + '_method') method_hosts = self.method_hosts_mapper[method_attr] method_hosts = [h for h in method_hosts if h != host['name']] inventory_hosts = [] records = [] host['secret_type'] = self.secret_type for account in accounts: h = deepcopy(host) h['name'] += '_' + account.username new_secret = self.get_secret() recorder = ChangeSecretRecord( asset=asset, account=account, execution=self.execution, old_secret=account.secret, new_secret=new_secret, ) records.append(recorder) self.name_recorder_mapper[h['name']] = recorder private_key_path = None if self.secret_type == SecretType.SSH_KEY: private_key_path = self.generate_private_key_path(new_secret, path_dir) new_secret = self.generate_public_key(new_secret) h['kwargs'] = self.get_kwargs(account, new_secret) h['account'] = { 'name': account.name, 'username': account.username, 'secret_type': account.secret_type, 'secret': new_secret, 'private_key_path': private_key_path } if asset.platform.type == 'oracle': h['account']['mode'] = 'sysdba' if account.privileged else None inventory_hosts.append(h) method_hosts.append(h['name']) self.method_hosts_mapper[method_attr] = method_hosts ChangeSecretRecord.objects.bulk_create(records) return inventory_hosts def on_host_success(self, host, result): recorder = self.name_recorder_mapper.get(host) if not recorder: return recorder.status = 'success' recorder.date_finished = timezone.now() recorder.save() account = recorder.account account.secret = recorder.new_secret account.save(update_fields=['secret']) def on_host_error(self, host, error, result): recorder = self.name_recorder_mapper.get(host) if not recorder: return recorder.status = 'failed' recorder.date_finished = timezone.now() recorder.error = error recorder.save() def on_runner_failed(self, runner, e): logger.error("Change secret error: ", e) def run(self, *args, **kwargs): super().run(*args, **kwargs) recorders = self.name_recorder_mapper.values() recorders = list(recorders) self.send_recorder_mail(recorders) def send_recorder_mail(self, recorders): recipients = self.execution.recipients if not recorders or not recipients: return recipients = User.objects.filter(id__in=list(recipients.keys())) name = self.execution.snapshot['name'] path = os.path.join(os.path.dirname(settings.BASE_DIR), 'tmp') filename = os.path.join(path, f'{name}-{local_now_display()}-{time.time()}.xlsx') if not self.create_file(recorders, filename): return for user in recipients: attachments = [] if user.secret_key: password = user.secret_key.encode('utf8') attachment = os.path.join(path, f'{name}-{local_now_display()}-{time.time()}.zip') encrypt_and_compress_zip_file(attachment, password, [filename]) attachments = [attachment] ChangeSecretExecutionTaskMsg(name, user).publish(attachments) os.remove(filename) @staticmethod def create_file(recorders, filename): serializer_cls = ChangeSecretRecordBackUpSerializer serializer = serializer_cls(recorders, many=True) header = [str(v.label) for v in serializer.child.fields.values()] rows = [list(row.values()) for row in serializer.data] if not rows: return False rows.insert(0, header) wb = Workbook(filename) ws = wb.create_sheet('Sheet1') for row in rows: ws.append(row) wb.save(filename) return True