perf: gathered-accounts api (#9451)

Co-authored-by: feng <1304903146@qq.com>

apps/accounts/api/account/__init__.py

@@ -1,3 +1,4 @@
 from .account import *
 from .backup import *
 from .template import *
+from .gathered_account import *

apps/accounts/api/account/gathered_account.py

@@ -0,0 +1,16 @@
+from accounts import serializers
+from accounts.models import GatheredAccount
+from orgs.mixins.api import OrgBulkModelViewSet
+
+__all__ = [
+    'GatheredAccountViewSet',
+]
+
+
+class GatheredAccountViewSet(OrgBulkModelViewSet):
+    model = GatheredAccount
+    search_fields = ('username', 'asset__address')
+    filterset_fields = ('username',)
+    serializer_classes = {
+        'default': serializers.AccountSerializer,
+    }

apps/accounts/automations/gather_accounts/manager.py

@@ -1,8 +1,7 @@
-from django.utils.translation import ugettext_lazy as _
-
-from common.utils import get_logger
-from accounts.const import AutomationTypes, Source
 from orgs.utils import tmp_to_org
+from common.utils import get_logger
+from accounts.models import GatheredAccount
+from accounts.const import AutomationTypes, Source
 from .filter import GatherAccountsFilter
 from ..base.manager import AccountBasePlaybookManager
 
@@ -28,32 +27,24 @@ class GatherAccountsManager(AccountBasePlaybookManager):
         return result
 
     @staticmethod
-    def bulk_create_accounts(asset, result):
-        account_objs = []
-        account_model = asset.accounts.model
-        account_usernames = set(asset.accounts.values_list('username', flat=True))
+    def update_or_create_gathered_accounts(asset, result):
         with tmp_to_org(asset.org_id):
-            accounts_dict = {}
+            GatheredAccount.objects.filter(asset=asset, present=True).update(present=False)
             for username, data in result.items():
-                comment = ''
-                d = {'asset': asset, 'username': username, 'name': username, 'source': Source.COLLECTED}
+                d = {'asset': asset, 'username': username, 'present': True}
                 if data.get('date'):
-                    comment += f"{_('Date last login')}: {data['date']}\n "
+                    d['date_last_login'] = data['date']
                 if data.get('address'):
-                    comment += f"{_('IP last login')}: {data['address'][:32]}"
-                d['comment'] = comment
-                accounts_dict[username] = d
-            for username, data in accounts_dict.items():
-                if username in account_usernames:
-                    continue
-                account_objs.append(account_model(**data))
-            account_model.objects.bulk_create(account_objs)
+                    d['address_last_login'] = data['address'][:32]
+                GatheredAccount.objects.update_or_create(
+                    defaults=d, asset=asset, username=username,
+                )
 
     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.bulk_create_accounts(asset, result)
+            self.update_or_create_gathered_accounts(asset, result)
         else:
             logger.error("Not found info".format(host))

apps/accounts/migrations/0006_gatheredaccount.py

@@ -0,0 +1,38 @@
+# Generated by Django 3.2.16 on 2023-02-07 04:41
+
+from django.db import migrations, models
+import django.db.models.deletion
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('assets', '0108_alter_platform_charset'),
+        ('accounts', '0005_alter_changesecretrecord_options'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='GatheredAccount',
+            fields=[
+                ('created_by', models.CharField(blank=True, max_length=128, null=True, verbose_name='Created by')),
+                ('updated_by', models.CharField(blank=True, max_length=128, null=True, verbose_name='Updated by')),
+                ('date_created', models.DateTimeField(auto_now_add=True, null=True, verbose_name='Date created')),
+                ('date_updated', models.DateTimeField(auto_now=True, verbose_name='Date updated')),
+                ('comment', models.TextField(blank=True, default='', verbose_name='Comment')),
+                ('id', models.UUIDField(default=uuid.uuid4, primary_key=True, serialize=False)),
+                ('org_id', models.CharField(blank=True, db_index=True, default='', max_length=36, verbose_name='Organization')),
+                ('present', models.BooleanField(default=True, verbose_name='Present')),
+                ('date_last_login', models.DateTimeField(null=True, verbose_name='Date last login')),
+                ('username', models.CharField(blank=True, db_index=True, max_length=32, verbose_name='Username')),
+                ('address_last_login', models.CharField(default='', max_length=39, verbose_name='Address last login')),
+                ('asset', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='assets.asset', verbose_name='Asset')),
+            ],
+            options={
+                'verbose_name': 'Gather account',
+                'ordering': ['asset'],
+                'unique_together': {('username', 'asset')},
+            },
+        ),
+    ]

apps/accounts/models/__init__.py

@@ -1,3 +1,4 @@
 from .base import *
 from .account import *
 from .automations import *
+from .gathered_account import *

apps/accounts/models/gathered_account.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+#
+from django.db import models
+from django.utils.translation import ugettext_lazy as _
+
+from orgs.mixins.models import JMSOrgBaseModel
+
+__all__ = ['GatheredAccount']
+
+
+class GatheredAccount(JMSOrgBaseModel):
+    present = models.BooleanField(default=True, verbose_name=_("Present"))
+    date_last_login = models.DateTimeField(null=True, verbose_name=_("Date last login"))
+    asset = models.ForeignKey('assets.Asset', on_delete=models.CASCADE, verbose_name=_("Asset"))
+    username = models.CharField(max_length=32, blank=True, db_index=True, verbose_name=_('Username'))
+    address_last_login = models.CharField(max_length=39, default='', verbose_name=_("Address last login"))
+
+    @property
+    def address(self):
+        return self.asset.address
+
+    class Meta:
+        verbose_name = _('Gather account')
+        unique_together = [
+            ('username', 'asset'),
+        ]
+        ordering = ['asset']
+
+    def __str__(self):
+        return '{}: {}'.format(self.asset, self.username)

apps/accounts/serializers/account/gathered_account.py

@@ -0,0 +1,23 @@
+from django.utils.translation import ugettext_lazy as _
+
+from accounts.models import GatheredAccount
+from orgs.mixins.serializers import BulkOrgResourceModelSerializer
+from .account import AccountAssetSerializer
+from .base import BaseAccountSerializer
+
+
+class AccountSerializer(BulkOrgResourceModelSerializer):
+    asset = AccountAssetSerializer(label=_('Asset'))
+
+    class Meta(BaseAccountSerializer.Meta):
+        model = GatheredAccount
+        fields = [
+            'present', 'asset', 'username',
+            'address_last_login', 'date_last_login'
+        ]
+
+    @classmethod
+    def setup_eager_loading(cls, queryset):
+        """ Perform necessary eager loading of data. """
+        queryset = queryset.prefetch_related('asset', 'asset__platform')
+        return queryset

apps/accounts/urls.py

@@ -9,6 +9,7 @@ app_name = 'accounts'
 router = BulkRouter()
 
 router.register(r'accounts', api.AccountViewSet, 'account')
+router.register(r'gathered-accounts', api.GatheredAccountViewSet, 'gathered-account')
 router.register(r'account-secrets', api.AccountSecretsViewSet, 'account-secret')
 router.register(r'account-templates', api.AccountTemplateViewSet, 'account-template')
 router.register(r'account-template-secrets', api.AccountTemplateSecretsViewSet, 'account-template-secret')

apps/assets/migrations/0108_alter_platform_charset.py

@@ -0,0 +1,18 @@
+# Generated by Django 3.2.16 on 2023-02-07 04:41
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('assets', '0107_automation'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='platform',
+            name='charset',
+            field=models.CharField(choices=[('utf-8', 'UTF-8'), ('gbk', 'GBK')], default='utf-8', max_length=8, verbose_name='Charset'),
+        ),
+    ]