feat: 收集账号 可选同步表

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)

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))

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'),
+        ),
+    ]

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)

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