From 1ac2fec13fbdf8f58676d79c6a55ceb17d4af6c2 Mon Sep 17 00:00:00 2001 From: feng <1304903146@qq.com> Date: Thu, 23 Mar 2023 18:57:22 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=B6=E9=9B=86=E8=B4=A6=E5=8F=B7=20?= =?UTF-8?q?=E5=8F=AF=E9=80=89=E5=90=8C=E6=AD=A5=E8=A1=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../api/automations/gather_accounts.py | 23 ++----------- .../automations/gather_accounts/manager.py | 33 ++++++++++++------ ...atheraccountsautomation_is_sync_account.py | 18 ++++++++++ .../models/automations/gather_account.py | 34 ++++++++++++++++++- .../automations/gather_accounts.py | 3 +- 5 files changed, 79 insertions(+), 32 deletions(-) create mode 100644 apps/accounts/migrations/0010_gatheraccountsautomation_is_sync_account.py diff --git a/apps/accounts/api/automations/gather_accounts.py b/apps/accounts/api/automations/gather_accounts.py index f4420f919..e6a846368 100644 --- a/apps/accounts/api/automations/gather_accounts.py +++ b/apps/accounts/api/automations/gather_accounts.py @@ -1,14 +1,13 @@ # -*- coding: utf-8 -*- # -from django.utils.translation import ugettext_lazy as _ from rest_framework import status from rest_framework.decorators import action from rest_framework.response import Response from accounts import serializers -from accounts.const import Source, AutomationTypes +from accounts.const import AutomationTypes from accounts.filters import GatheredAccountFilterSet -from accounts.models import GatherAccountsAutomation, Account +from accounts.models import GatherAccountsAutomation from accounts.models import GatheredAccount from orgs.mixins.api import OrgBulkModelViewSet from .base import AutomationExecutionViewSet @@ -56,21 +55,5 @@ class GatheredAccountViewSet(OrgBulkModelViewSet): def sync_accounts(self, request, *args, **kwargs): gathered_account_ids = request.data.get('gathered_account_ids') gathered_accounts = self.model.objects.filter(id__in=gathered_account_ids) - account_objs = [] - exists_accounts = Account.objects.none() - for gathered_account in gathered_accounts: - asset_id = gathered_account.asset_id - username = gathered_account.username - accounts = Account.objects.filter(asset_id=asset_id, username=username) - if accounts.exists(): - exists_accounts |= accounts - else: - account_objs.append( - Account( - asset_id=asset_id, username=username, - name=f'{username}-{_("Collected")}', - source=Source.COLLECTED - )) - exists_accounts.update(source=Source.COLLECTED) - Account.objects.bulk_create(account_objs) + self.model.sync_accounts(gathered_accounts) return Response(status=status.HTTP_201_CREATED) diff --git a/apps/accounts/automations/gather_accounts/manager.py b/apps/accounts/automations/gather_accounts/manager.py index 2ecd3d2e1..c4ba6b5a0 100644 --- a/apps/accounts/automations/gather_accounts/manager.py +++ b/apps/accounts/automations/gather_accounts/manager.py @@ -12,6 +12,7 @@ class GatherAccountsManager(AccountBasePlaybookManager): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.host_asset_mapper = {} + self.is_sync_account = self.execution.snapshot.get('is_sync_account') @classmethod def method_type(cls): @@ -25,26 +26,38 @@ class GatherAccountsManager(AccountBasePlaybookManager): def filter_success_result(self, tp, result): result = GatherAccountsFilter(tp).run(self.method_id_meta_mapper, result) return result - @staticmethod - def update_or_create_gathered_accounts(asset, result): + def generate_data(asset, result): + data = [] + for username, info in result.items(): + d = {'asset': asset, 'username': username, 'present': True} + if info.get('date'): + d['date_last_login'] = info['date'] + if info.get('address'): + d['address_last_login'] = info['address'][:32] + data.append(d) + return data + + def update_or_create_accounts(self, asset, result): + data = self.generate_data(asset, result) with tmp_to_org(asset.org_id): + gathered_accounts = [] GatheredAccount.objects.filter(asset=asset, present=True).update(present=False) - for username, data in result.items(): - d = {'asset': asset, 'username': username, 'present': True} - if data.get('date'): - d['date_last_login'] = data['date'] - if data.get('address'): - d['address_last_login'] = data['address'][:32] - GatheredAccount.objects.update_or_create( + for d in data: + username = d['username'] + gathered_account, __ = GatheredAccount.objects.update_or_create( defaults=d, asset=asset, username=username, ) + gathered_accounts.append(gathered_account) + if not self.is_sync_account: + return + GatheredAccount.sync_accounts(gathered_accounts) def on_host_success(self, host, result): info = result.get('debug', {}).get('res', {}).get('info', {}) asset = self.host_asset_mapper.get(host) if asset and info: result = self.filter_success_result(asset.type, info) - self.update_or_create_gathered_accounts(asset, result) + self.update_or_create_accounts(asset, result) else: logger.error("Not found info".format(host)) diff --git a/apps/accounts/migrations/0010_gatheraccountsautomation_is_sync_account.py b/apps/accounts/migrations/0010_gatheraccountsautomation_is_sync_account.py new file mode 100644 index 000000000..058a7dd7e --- /dev/null +++ b/apps/accounts/migrations/0010_gatheraccountsautomation_is_sync_account.py @@ -0,0 +1,18 @@ +# Generated by Django 3.2.16 on 2023-03-23 08:39 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('accounts', '0009_account_usernames_to_ids'), + ] + + operations = [ + migrations.AddField( + model_name='gatheraccountsautomation', + name='is_sync_account', + field=models.BooleanField(blank=True, default=False, verbose_name='Is sync account'), + ), + ] diff --git a/apps/accounts/models/automations/gather_account.py b/apps/accounts/models/automations/gather_account.py index 01f903141..dd9b5c862 100644 --- a/apps/accounts/models/automations/gather_account.py +++ b/apps/accounts/models/automations/gather_account.py @@ -1,7 +1,9 @@ from django.db import models +from django.db.models import Q from django.utils.translation import ugettext_lazy as _ -from accounts.const import AutomationTypes +from accounts.const import AutomationTypes, Source +from accounts.models import Account from orgs.mixins.models import JMSOrgBaseModel from .base import AccountBaseAutomation @@ -19,6 +21,25 @@ class GatheredAccount(JMSOrgBaseModel): def address(self): return self.asset.address + @staticmethod + def sync_accounts(gathered_accounts): + account_objs = [] + for gathered_account in gathered_accounts: + asset_id = gathered_account.asset_id + username = gathered_account.username + accounts = Account.objects.filter( + Q(asset_id=asset_id, username=username) | + Q(asset_id=asset_id, name=username) + ) + if accounts.exists(): + continue + account = Account( + asset_id=asset_id, username=username, + name=username, source=Source.COLLECTED + ) + account_objs.append(account) + Account.objects.bulk_create(account_objs) + class Meta: verbose_name = _('Gather account automation') unique_together = [ @@ -31,6 +52,17 @@ class GatheredAccount(JMSOrgBaseModel): class GatherAccountsAutomation(AccountBaseAutomation): + is_sync_account = models.BooleanField( + default=False, blank=True, verbose_name=_("Is sync account") + ) + + def to_attr_json(self): + attr_json = super().to_attr_json() + attr_json.update({ + 'is_sync_account': self.is_sync_account, + }) + return attr_json + def save(self, *args, **kwargs): self.type = AutomationTypes.gather_accounts super().save(*args, **kwargs) diff --git a/apps/accounts/serializers/automations/gather_accounts.py b/apps/accounts/serializers/automations/gather_accounts.py index b906e7881..fd09773de 100644 --- a/apps/accounts/serializers/automations/gather_accounts.py +++ b/apps/accounts/serializers/automations/gather_accounts.py @@ -17,7 +17,8 @@ class GatherAccountAutomationSerializer(BaseAutomationSerializer): class Meta: model = GatherAccountsAutomation read_only_fields = BaseAutomationSerializer.Meta.read_only_fields - fields = BaseAutomationSerializer.Meta.fields + read_only_fields + fields = BaseAutomationSerializer.Meta.fields \ + + ['is_sync_account'] + read_only_fields extra_kwargs = BaseAutomationSerializer.Meta.extra_kwargs