perf: 修改 account remote

apps/accounts/automations/remove_account/manager.py

@@ -6,6 +6,7 @@ from django.db.models import QuerySet
 
 from accounts.const import AutomationTypes
 from accounts.models import Account, GatheredAccount, AccountRisk
+from common.const import ConfirmOrIgnore
 from common.utils import get_logger
 from ..base.manager import AccountBasePlaybookManager
 
@@ -24,6 +25,10 @@ class RemoveAccountManager(AccountBasePlaybookManager):
         for account in snapshot_account:
             self.snapshot_asset_account_map[str(account["asset"])].append(account)
 
+        # 给 handler 使用
+        self.delete = self.execution.snapshot.get("delete", "both")
+        self.confirm_risk = self.execution.snapshot.get("risk", "")
+
     def prepare_runtime_dir(self):
         path = super().prepare_runtime_dir()
         ansible_config_path = os.path.join(path, "ansible.cfg")
@@ -66,18 +71,24 @@ class RemoveAccountManager(AccountBasePlaybookManager):
             return
 
         try:
-            Account.objects.filter(
-                asset_id=account["asset"], username=account["username"]
-            ).delete()
+            if self.delete == "both":
+                Account.objects.filter(
+                    asset_id=account["asset"],
+                    username=account["username"]
+                ).delete()
+
+            if self.confirm_risk:
+                AccountRisk.objects.filter(
+                    asset_id=account["asset"],
+                    username=account["username"],
+                    risk__in=[self.confirm_risk],
+                ).update(status=ConfirmOrIgnore.confirmed)
+
             GatheredAccount.objects.filter(
-                asset_id=account["asset"], username=account["username"]
-            ).delete()
-            risk = AccountRisk.objects.filter(
                 asset_id=account["asset"],
-                username=account["username"],
-                risk__in=["new_found"],
-            )
-            print("Account removed: ", account)
+                username=account["username"]
+            ).delete()
+
         except Exception as e:
             logger.error(
                 f"Failed to delete account {account['username']} on asset {account['asset']}: {e}"

apps/accounts/risk_handlers.py

@@ -1,7 +1,13 @@
 from django.utils.translation import gettext_lazy as _
 
+from accounts.const import AutomationTypes
 from common.const import ConfirmOrIgnore
-from accounts.models import GatheredAccount, AccountRisk, SecretType, AutomationExecution
+from accounts.models import (
+    GatheredAccount,
+    AccountRisk,
+    SecretType,
+    AutomationExecution,
+)
 from django.utils import timezone
 
 from common.const import ConfirmOrIgnore
@@ -18,13 +24,13 @@ TYPE_CHOICES = [
 
 
 class RiskHandler:
-    def __init__(self, asset, username, request=None, risk=''):
+    def __init__(self, asset, username, request=None, risk=""):
         self.asset = asset
         self.username = username
         self.request = request
         self.risk = risk
 
-    def handle(self, tp, risk=''):
+    def handle(self, tp, risk=""):
         self.risk = risk
         attr = f"handle_{tp}"
         if hasattr(self, attr):
@@ -38,11 +44,10 @@ class RiskHandler:
         r = self.get_risk()
         if not r:
             return
-        status = ConfirmOrIgnore.ignored if tp == 'ignore' else ConfirmOrIgnore.confirmed
-        r.details.append({
-            **self.process_detail,
-            'action': tp, 'status': status
-        })
+        status = (
+            ConfirmOrIgnore.ignored if tp == "ignore" else ConfirmOrIgnore.confirmed
+        )
+        r.details.append({**self.process_detail, "action": tp, "status": status})
         r.status = status
         r.save()
 
@@ -61,8 +66,9 @@ class RiskHandler:
     @property
     def process_detail(self):
         return {
-            "datetime": timezone.now().isoformat(), "type": "process",
-            "processor": str(self.request.user)
+            "datetime": timezone.now().isoformat(),
+            "type": "process",
+            "processor": str(self.request.user),
         }
 
     def handle_add_account(self):
@@ -76,12 +82,15 @@ class RiskHandler:
         GatheredAccount.objects.filter(asset=self.asset, username=self.username).update(
             present=True, status=ConfirmOrIgnore.confirmed
         )
-        self.risk = 'new_found'
+        self.risk = "new_found"
 
     def handle_disable_remote(self):
         pass
 
     def handle_delete_remote(self):
+        self._handle_delete(delete="remote")
+
+    def _handle_delete(self, delete="both"):
         asset = self.asset
         execution = AutomationExecution()
         execution.snapshot = {
@@ -89,16 +98,30 @@ class RiskHandler:
             "accounts": [{"asset": str(asset.id), "username": self.username}],
             "type": "remove_account",
             "name": "Remove remote account: {}@{}".format(self.username, asset.name),
+            "delete": delete,
+            "risk": self.risk
         }
         execution.save()
         execution.start()
         return execution.summary
 
     def handle_delete_both(self):
-        pass
+        self._handle_delete(delete="both")
 
     def handle_change_password_add(self):
         pass
 
     def handle_change_password(self):
-        pass
+        asset = self.asset
+        execution = AutomationExecution()
+        execution.snapshot = {
+            "assets": [str(asset.id)],
+            "accounts": [self.username],
+            "type": AutomationTypes.change_secret,
+            "secret_type": "password",
+            "secret_strategy": "random",
+            "name": "Change account password: {}@{}".format(self.username, asset.name),
+        }
+        execution.save()
+        execution.start()
+        return execution.summary