From eefda353d2da0ab6c93066414c079d56e69fd270 Mon Sep 17 00:00:00 2001 From: feng <1304903146@qq.com> Date: Fri, 6 Dec 2024 11:30:26 +0800 Subject: [PATCH] perf: Change secret --- .../automations/change_secret/manager.py | 52 +++++++++---------- .../automations/push_account/manager.py | 3 ++ apps/accounts/models/automations/base.py | 18 ++++--- .../models/automations/push_account.py | 10 ++-- apps/accounts/models/base.py | 5 +- apps/accounts/models/template.py | 4 ++ apps/accounts/serializers/account/template.py | 17 ------ 7 files changed, 53 insertions(+), 56 deletions(-) diff --git a/apps/accounts/automations/change_secret/manager.py b/apps/accounts/automations/change_secret/manager.py index 285c4ba04..01078a057 100644 --- a/apps/accounts/automations/change_secret/manager.py +++ b/apps/accounts/automations/change_secret/manager.py @@ -46,7 +46,7 @@ class ChangeSecretManager(AccountBasePlaybookManager): def method_type(cls): return AutomationTypes.change_secret - def get_ssh_params(self, account, secret, secret_type): + def get_ssh_params(self, secret, secret_type): kwargs = {} if secret_type != SecretType.SSH_KEY: return kwargs @@ -65,6 +65,7 @@ class ChangeSecretManager(AccountBasePlaybookManager): asset = privilege_account.asset accounts = asset.accounts.all() accounts = accounts.filter(id__in=self.account_ids) + if self.secret_type: accounts = accounts.filter(secret_type=self.secret_type) @@ -74,37 +75,36 @@ class ChangeSecretManager(AccountBasePlaybookManager): ) return accounts - def gen_new_secret(self, account, path_dir): - if self.secret_type is None: - new_secret = account.secret + def get_secret(self, account): + if self.secret_strategy == SecretStrategy.custom: + new_secret = self.execution.snapshot['secret'] else: - if self.secret_strategy == SecretStrategy.custom: - new_secret = self.execution.snapshot['secret'] - else: - generator = SecretGenerator( - self.secret_strategy, self.secret_type, - self.execution.snapshot.get('password_rules') - ) - new_secret = generator.get_secret() + generator = SecretGenerator( + self.secret_strategy, self.secret_type, + self.execution.snapshot.get('password_rules') + ) + new_secret = generator.get_secret() + return new_secret - recorder_secret = new_secret + def gen_new_secret(self, account, new_secret, path_dir): private_key_path = None if account.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) - return new_secret, private_key_path, recorder_secret + return new_secret, private_key_path - def get_or_create_record(self, asset, account, new_secret, name): + def get_or_create_record(self, asset, account, name): asset_account_id = f'{asset.id}-{account.id}' if asset_account_id in self.record_map: record_id = self.record_map[asset_account_id] recorder = ChangeSecretRecord.objects.filter(id=record_id).first() else: + new_secret = self.get_secret(account) recorder = self.create_record(asset, account, new_secret) - if recorder: - self.name_recorder_mapper[name] = recorder + self.name_recorder_mapper[name] = recorder + return recorder @bulk_create_decorator(ChangeSecretRecord) def create_record(self, asset, account, new_secret): @@ -115,11 +115,9 @@ class ChangeSecretManager(AccountBasePlaybookManager): ) return recorder - def gen_change_secret_inventory(self, host, account, new_secret, private_key_path, asset): - h = deepcopy(host) + def gen_change_secret_inventory(self, h, account, new_secret, private_key_path, asset): secret_type = account.secret_type - h['name'] += '(' + account.username + ')' - h['ssh_params'].update(self.get_ssh_params(account, new_secret, secret_type)) + h['ssh_params'].update(self.get_ssh_params(new_secret, secret_type)) h['account'] = { 'name': account.name, 'username': account.username, @@ -144,7 +142,7 @@ class ChangeSecretManager(AccountBasePlaybookManager): host['ssh_params'] = {} accounts = self.get_accounts(account) - error_msg = _("! No pending accounts found") + error_msg = _("No pending accounts found") if not accounts: print(f'{asset}: {error_msg}') return [] @@ -154,13 +152,15 @@ class ChangeSecretManager(AccountBasePlaybookManager): inventory_hosts = [] if asset.type == HostTypes.WINDOWS and self.secret_type == SecretType.SSH_KEY: - print(f'! Windows {asset} does not support ssh key push') + print(f'Windows {asset} does not support ssh key push') return inventory_hosts for account in accounts: - new_secret, private_key_path, recorder_secret = self.gen_new_secret(account, path_dir) - h = self.gen_change_secret_inventory(host, account, new_secret, private_key_path, asset) - self.get_or_create_record(asset, account, recorder_secret, h['name']) + h = deepcopy(host) + h['name'] += '(' + account.username + ')' # To distinguish different accounts + record = self.get_or_create_record(asset, account, h['name']) + new_secret, private_key_path = self.gen_new_secret(account, record.new_secret, path_dir) + h = self.gen_change_secret_inventory(h, account, new_secret, private_key_path, asset) inventory_hosts.append(h) return inventory_hosts diff --git a/apps/accounts/automations/push_account/manager.py b/apps/accounts/automations/push_account/manager.py index 18aca5134..91463f96b 100644 --- a/apps/accounts/automations/push_account/manager.py +++ b/apps/accounts/automations/push_account/manager.py @@ -10,3 +10,6 @@ class PushAccountManager(ChangeSecretManager, AccountBasePlaybookManager): @classmethod def method_type(cls): return AutomationTypes.push_account + + def get_secret(self, account): + return account.secret diff --git a/apps/accounts/models/automations/base.py b/apps/accounts/models/automations/base.py index 799cbd3d7..19fa891e4 100644 --- a/apps/accounts/models/automations/base.py +++ b/apps/accounts/models/automations/base.py @@ -64,19 +64,25 @@ class ChangeSecretMixin(SecretWithRandomMixin): verbose_name=_('Check connection after change') ) get_all_assets: callable # get all assets + accounts: list # account usernames class Meta: abstract = True - def create_nonlocal_accounts(self, usernames, asset): - pass + def gen_nonlocal_accounts(self, usernames, asset): + return [] def get_account_ids(self): + account_objs = [] usernames = self.accounts - accounts = Account.objects.none() - for asset in self.get_all_assets(): - self.create_nonlocal_accounts(usernames, asset) - accounts = accounts | asset.accounts.all() + assets = self.get_all_assets() + for asset in assets: + objs = self.gen_nonlocal_accounts(usernames, asset) + account_objs.extend(objs) + + Account.objects.bulk_create(account_objs) + + accounts = Account.objects.filter(asset__in=assets) account_ids = accounts.filter( username__in=usernames, secret_type=self.secret_type ).values_list('id', flat=True) diff --git a/apps/accounts/models/automations/push_account.py b/apps/accounts/models/automations/push_account.py index d3eb197a4..5a5ee2d51 100644 --- a/apps/accounts/models/automations/push_account.py +++ b/apps/accounts/models/automations/push_account.py @@ -11,20 +11,22 @@ __all__ = ['PushAccountAutomation'] class PushAccountAutomation(ChangeSecretMixin, AccountBaseAutomation): - def create_nonlocal_accounts(self, usernames, asset): + def gen_nonlocal_accounts(self, usernames, asset): secret_type = self.secret_type account_usernames = asset.accounts \ .filter(secret_type=self.secret_type) \ .values_list('username', flat=True) create_usernames = set(usernames) - set(account_usernames) - create_account_objs = [ + + create_accounts = [ Account( - name=f'{username}-{secret_type}', username=username, + name=f'{username}-{secret_type}', + username=username, secret=self.get_secret(), secret_type=secret_type, asset=asset, ) for username in create_usernames ] - Account.objects.bulk_create(create_account_objs) + return create_accounts def save(self, *args, **kwargs): self.type = AutomationTypes.push_account diff --git a/apps/accounts/models/base.py b/apps/accounts/models/base.py index 4b55752d6..583ebd318 100644 --- a/apps/accounts/models/base.py +++ b/apps/accounts/models/base.py @@ -54,10 +54,9 @@ class SecretWithRandomMixin(models.Model): ) def get_secret(self): - if self.secret_strategy == 'random': - return self.secret_generator.get_secret() - else: + if self.secret_strategy == SecretStrategy.custom: return self.secret + return self.secret_generator.get_secret() class BaseAccount(VaultModelMixin, JMSOrgBaseModel): diff --git a/apps/accounts/models/template.py b/apps/accounts/models/template.py index 8d5be852a..6a24caf85 100644 --- a/apps/accounts/models/template.py +++ b/apps/accounts/models/template.py @@ -86,3 +86,7 @@ class AccountTemplate(LabeledMixin, BaseAccount, SecretWithRandomMixin): """ 批量同步账号密码 """ self.bulk_update_accounts(accounts) self.bulk_create_history_accounts(accounts, user_id) + + def save(self, *args, **kwargs): + self.secret = self.get_secret() + super().save(*args, **kwargs) diff --git a/apps/accounts/serializers/account/template.py b/apps/accounts/serializers/account/template.py index fc6ccefeb..b5ab0e00e 100644 --- a/apps/accounts/serializers/account/template.py +++ b/apps/accounts/serializers/account/template.py @@ -1,9 +1,7 @@ from django.utils.translation import gettext_lazy as _ from rest_framework import serializers -from accounts.const import SecretStrategy, SecretType from accounts.models import AccountTemplate -from accounts.utils import SecretGenerator from common.serializers import SecretReadableMixin from common.serializers.fields import ObjectRelatedField from .base import BaseAccountSerializer @@ -47,21 +45,6 @@ class AccountTemplateSerializer(BaseAccountSerializer): }, } - @staticmethod - def generate_secret(attrs): - secret_type = attrs.get('secret_type', SecretType.PASSWORD) - secret_strategy = attrs.get('secret_strategy', SecretStrategy.custom) - password_rules = attrs.get('password_rules') - if secret_strategy != SecretStrategy.random: - return - generator = SecretGenerator(secret_strategy, secret_type, password_rules) - attrs['secret'] = generator.get_secret() - - def validate(self, attrs): - attrs = super().validate(attrs) - self.generate_secret(attrs) - return attrs - class AccountTemplateSecretSerializer(SecretReadableMixin, AccountTemplateSerializer): class Meta(AccountTemplateSerializer.Meta):