perf: risk add account

pull/14806/head
ibuler 2025-01-13 11:14:26 +08:00
parent f7d8e7952a
commit 860af61557
5 changed files with 53 additions and 21 deletions

View File

@ -221,6 +221,24 @@ class CheckAccountManager(BaseManager):
self.execution.date_start = timezone.now() self.execution.date_start = timezone.now()
self.execution.save(update_fields=["date_start"]) self.execution.save(update_fields=["date_start"])
def batch_check(self, handler):
print("Engine: {}".format(handler.__class__.__name__))
for i in range(0, len(self.assets), self.batch_size):
_assets = self.assets[i: i + self.batch_size]
accounts = Account.objects.filter(asset__in=_assets)
print("Start to check accounts: {}".format(len(accounts)))
for account in accounts:
error = handler.check(account)
msg = handler.risk if error else 'ok'
print("Check: {} => {}".format(account, msg))
if not error:
continue
self.add_risk(handler.risk, account)
self.commit_risks(_assets)
def do_run(self, *args, **kwargs): def do_run(self, *args, **kwargs):
for engine in self.execution.snapshot.get("engines", []): for engine in self.execution.snapshot.get("engines", []):
if engine == "check_account_secret": if engine == "check_account_secret":
@ -234,22 +252,7 @@ class CheckAccountManager(BaseManager):
continue continue
self.handlers.append(handler) self.handlers.append(handler)
self.batch_check(handler)
print("Engine: {}".format(handler.__class__.__name__))
for i in range(0, len(self.assets), self.batch_size):
_assets = self.assets[i: i + self.batch_size]
accounts = Account.objects.filter(asset__in=_assets)
print("Start to check accounts: {}".format(len(accounts)))
for account in accounts:
error = handler.check(account)
print("Check: {} => {}".format(account, error))
if not error:
continue
self.add_risk(handler.risk, account)
self.commit_risks(_assets)
def post_run(self): def post_run(self):
super().post_run() super().post_run()

View File

@ -0,0 +1,25 @@
# Generated by Django 4.1.13 on 2025-01-13 03:13
from django.db import migrations, models
import django.db.models.deletion
class Migration(migrations.Migration):
dependencies = [
("accounts", "0025_alter_accountrisk_risk_and_more"),
]
operations = [
migrations.AddField(
model_name="accountrisk",
name="account",
field=models.ForeignKey(
null=True,
on_delete=django.db.models.deletion.CASCADE,
related_name="risks",
to="accounts.account",
verbose_name="Account",
),
),
]

View File

@ -2,8 +2,8 @@ from itertools import islice
from django.db import models from django.db import models
from django.db.models import TextChoices from django.db.models import TextChoices
from django.utils.translation import gettext_lazy as _
from django.utils import timezone from django.utils import timezone
from django.utils.translation import gettext_lazy as _
from common.const import ConfirmOrIgnore from common.const import ConfirmOrIgnore
from common.db.models import JMSBaseModel from common.db.models import JMSBaseModel
@ -60,9 +60,12 @@ class RiskChoice(TextChoices):
class AccountRisk(JMSOrgBaseModel): class AccountRisk(JMSOrgBaseModel):
asset = models.ForeignKey('assets.Asset', on_delete=models.CASCADE, related_name='risks', verbose_name=_('Asset')) asset = models.ForeignKey('assets.Asset', on_delete=models.CASCADE, related_name='risks', verbose_name=_('Asset'))
username = models.CharField(max_length=32, verbose_name=_('Username')) username = models.CharField(max_length=32, verbose_name=_('Username'))
account = models.ForeignKey('accounts.Account', on_delete=models.CASCADE, related_name='risks',
verbose_name=_('Account'), null=True)
risk = models.CharField(max_length=128, verbose_name=_('Risk'), choices=RiskChoice.choices) risk = models.CharField(max_length=128, verbose_name=_('Risk'), choices=RiskChoice.choices)
status = models.CharField(max_length=32, choices=ConfirmOrIgnore.choices, default=ConfirmOrIgnore.pending, blank=True, verbose_name=_('Status')) status = models.CharField(max_length=32, choices=ConfirmOrIgnore.choices, default=ConfirmOrIgnore.pending,
details = models.JSONField(default=list, verbose_name=_('Details')) blank=True, verbose_name=_('Status'))
details = models.JSONField(default=list, verbose_name=_('Details'))
class Meta: class Meta:
verbose_name = _('Account risk') verbose_name = _('Account risk')
@ -106,7 +109,7 @@ class AccountRisk(JMSOrgBaseModel):
class CheckAccountEngine(JMSBaseModel): class CheckAccountEngine(JMSBaseModel):
name = models.CharField(max_length=128, verbose_name=_('Name'), unique=True) name = models.CharField(max_length=128, verbose_name=_('Name'), unique=True)
slug = models.SlugField(max_length=128, verbose_name=_('Slug'), unique=True) # slug = models.SlugField(max_length=128, verbose_name=_('Slug'), unique=True) #
is_active = models.BooleanField(default=True, verbose_name=_('Is active')) is_active = models.BooleanField(default=True, verbose_name=_('Is active'))
def __str__(self): def __str__(self):
@ -117,4 +120,3 @@ class CheckAccountEngine(JMSBaseModel):
'check_gathered_account', 'check_gathered_account',
'check_account_secret' 'check_account_secret'
] ]

View File

@ -996,6 +996,7 @@
"ResetSSHKeyWarningMsg": "Are you sure you want to send a reset ssh key email to the user?", "ResetSSHKeyWarningMsg": "Are you sure you want to send a reset ssh key email to the user?",
"Resource": "Resources", "Resource": "Resources",
"ResourceType": "Resource type", "ResourceType": "Resource type",
"ResolveSelected": "Resolve selected",
"RestoreButton": "Restore", "RestoreButton": "Restore",
"RestoreDefault": "Reset to default", "RestoreDefault": "Reset to default",
"RestoreDialogMessage": "Are you sure you want to restore to default initialization?", "RestoreDialogMessage": "Are you sure you want to restore to default initialization?",

View File

@ -994,6 +994,7 @@
"ResetSSHKeyWarningMsg": "你确定要发送重置用户的SSH Key的邮件吗?", "ResetSSHKeyWarningMsg": "你确定要发送重置用户的SSH Key的邮件吗?",
"Resource": "资源", "Resource": "资源",
"ResourceType": "资源类型", "ResourceType": "资源类型",
"ResolveSelected": "解决所选",
"RestoreButton": "恢复默认", "RestoreButton": "恢复默认",
"RestoreDefault": "恢复默认", "RestoreDefault": "恢复默认",
"RestoreDialogMessage": "您确定要恢复默认初始化吗?", "RestoreDialogMessage": "您确定要恢复默认初始化吗?",