From a9f9e46a37a3cc7b8403b27eacea18d62a84f490 Mon Sep 17 00:00:00 2001 From: ibuler Date: Mon, 9 Dec 2024 17:11:03 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BF=AE=E6=94=B9=20account=20remote?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../automations/remove_account/manager.py | 31 ++++++++---- apps/accounts/risk_handlers.py | 49 ++++++++++++++----- 2 files changed, 57 insertions(+), 23 deletions(-) diff --git a/apps/accounts/automations/remove_account/manager.py b/apps/accounts/automations/remove_account/manager.py index fb95f15cc..2d194164f 100644 --- a/apps/accounts/automations/remove_account/manager.py +++ b/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}" diff --git a/apps/accounts/risk_handlers.py b/apps/accounts/risk_handlers.py index c7161a969..389b8d963 100644 --- a/apps/accounts/risk_handlers.py +++ b/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