From 376a9d74e82b6882e93cd780061e32711bc0cfa0 Mon Sep 17 00:00:00 2001
From: feng <1304903146@qq.com>
Date: Mon, 17 Mar 2025 18:43:46 +0800
Subject: [PATCH] perf: push account secret

---
 apps/accounts/automations/base/manager.py          | 12 ++++++++++++
 apps/accounts/automations/change_secret/manager.py | 12 ------------
 apps/accounts/automations/push_account/manager.py  |  5 ++++-
 3 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/apps/accounts/automations/base/manager.py b/apps/accounts/automations/base/manager.py
index f33ea358f..1782d0f87 100644
--- a/apps/accounts/automations/base/manager.py
+++ b/apps/accounts/automations/base/manager.py
@@ -7,6 +7,7 @@ from django.utils.translation import gettext_lazy as _
 from accounts.automations.methods import platform_automation_methods
 from accounts.const import SSHKeyStrategy, SecretStrategy, SecretType, ChangeSecretRecordStatusChoice
 from accounts.models import BaseAccountQuerySet
+from accounts.utils import SecretGenerator
 from assets.automations.base.manager import BasePlaybookManager
 from assets.const import HostTypes
 from common.db.utils import safe_db_connection
@@ -51,6 +52,17 @@ class BaseChangeSecretPushManager(AccountBasePlaybookManager):
             kwargs['regexp'] = '.*{}$'.format(secret.split()[2].strip())
         return kwargs
 
+    def get_secret(self, account):
+        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()
+        return new_secret
+
     def get_accounts(self, privilege_account) -> BaseAccountQuerySet | None:
         if not privilege_account:
             print('Not privilege account')
diff --git a/apps/accounts/automations/change_secret/manager.py b/apps/accounts/automations/change_secret/manager.py
index 2c56b1e73..4cbec576d 100644
--- a/apps/accounts/automations/change_secret/manager.py
+++ b/apps/accounts/automations/change_secret/manager.py
@@ -15,7 +15,6 @@ from common.utils import get_logger
 from common.utils.file import encrypt_and_compress_zip_file
 from common.utils.timezone import local_now_filename
 from ..base.manager import BaseChangeSecretPushManager
-from ...utils import SecretGenerator
 
 logger = get_logger(__name__)
 
@@ -27,17 +26,6 @@ class ChangeSecretManager(BaseChangeSecretPushManager):
     def method_type(cls):
         return AutomationTypes.change_secret
 
-    def get_secret(self, account):
-        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()
-        return new_secret
-
     def gen_account_inventory(self, account, asset, h, path_dir):
         record = self.get_or_create_record(asset, account, h['name'])
         new_secret, private_key_path = self.handle_ssh_secret(account.secret_type, record.new_secret, path_dir)
diff --git a/apps/accounts/automations/push_account/manager.py b/apps/accounts/automations/push_account/manager.py
index 2ba85a527..6975f5ec6 100644
--- a/apps/accounts/automations/push_account/manager.py
+++ b/apps/accounts/automations/push_account/manager.py
@@ -21,7 +21,10 @@ class PushAccountManager(BaseChangeSecretPushManager):
         return AutomationTypes.push_account
 
     def get_secret(self, account):
-        return account.secret
+        secret = account.secret
+        if not secret:
+            secret = super().get_secret(account)
+        return secret
 
     def gen_account_inventory(self, account, asset, h, path_dir):
         secret = self.get_secret(account)