perf: Change secret

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

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

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)

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

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

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)

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