feat: 收集账号 可选同步表

pull/10075/head
feng 2023-03-23 18:57:22 +08:00 committed by Jiangjie.Bai
parent 66d368f882
commit 1ac2fec13f
5 changed files with 79 additions and 32 deletions

View File

@ -1,14 +1,13 @@
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
# #
from django.utils.translation import ugettext_lazy as _
from rest_framework import status from rest_framework import status
from rest_framework.decorators import action from rest_framework.decorators import action
from rest_framework.response import Response from rest_framework.response import Response
from accounts import serializers from accounts import serializers
from accounts.const import Source, AutomationTypes from accounts.const import AutomationTypes
from accounts.filters import GatheredAccountFilterSet from accounts.filters import GatheredAccountFilterSet
from accounts.models import GatherAccountsAutomation, Account from accounts.models import GatherAccountsAutomation
from accounts.models import GatheredAccount from accounts.models import GatheredAccount
from orgs.mixins.api import OrgBulkModelViewSet from orgs.mixins.api import OrgBulkModelViewSet
from .base import AutomationExecutionViewSet from .base import AutomationExecutionViewSet
@ -56,21 +55,5 @@ class GatheredAccountViewSet(OrgBulkModelViewSet):
def sync_accounts(self, request, *args, **kwargs): def sync_accounts(self, request, *args, **kwargs):
gathered_account_ids = request.data.get('gathered_account_ids') gathered_account_ids = request.data.get('gathered_account_ids')
gathered_accounts = self.model.objects.filter(id__in=gathered_account_ids) gathered_accounts = self.model.objects.filter(id__in=gathered_account_ids)
account_objs = [] self.model.sync_accounts(gathered_accounts)
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)
return Response(status=status.HTTP_201_CREATED) return Response(status=status.HTTP_201_CREATED)

View File

@ -12,6 +12,7 @@ class GatherAccountsManager(AccountBasePlaybookManager):
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
self.host_asset_mapper = {} self.host_asset_mapper = {}
self.is_sync_account = self.execution.snapshot.get('is_sync_account')
@classmethod @classmethod
def method_type(cls): def method_type(cls):
@ -25,26 +26,38 @@ class GatherAccountsManager(AccountBasePlaybookManager):
def filter_success_result(self, tp, result): def filter_success_result(self, tp, result):
result = GatherAccountsFilter(tp).run(self.method_id_meta_mapper, result) result = GatherAccountsFilter(tp).run(self.method_id_meta_mapper, result)
return result return result
@staticmethod @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): with tmp_to_org(asset.org_id):
gathered_accounts = []
GatheredAccount.objects.filter(asset=asset, present=True).update(present=False) GatheredAccount.objects.filter(asset=asset, present=True).update(present=False)
for username, data in result.items(): for d in data:
d = {'asset': asset, 'username': username, 'present': True} username = d['username']
if data.get('date'): gathered_account, __ = GatheredAccount.objects.update_or_create(
d['date_last_login'] = data['date']
if data.get('address'):
d['address_last_login'] = data['address'][:32]
GatheredAccount.objects.update_or_create(
defaults=d, asset=asset, username=username, 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): def on_host_success(self, host, result):
info = result.get('debug', {}).get('res', {}).get('info', {}) info = result.get('debug', {}).get('res', {}).get('info', {})
asset = self.host_asset_mapper.get(host) asset = self.host_asset_mapper.get(host)
if asset and info: if asset and info:
result = self.filter_success_result(asset.type, 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: else:
logger.error("Not found info".format(host)) logger.error("Not found info".format(host))

View File

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

View File

@ -1,7 +1,9 @@
from django.db import models from django.db import models
from django.db.models import Q
from django.utils.translation import ugettext_lazy as _ 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 orgs.mixins.models import JMSOrgBaseModel
from .base import AccountBaseAutomation from .base import AccountBaseAutomation
@ -19,6 +21,25 @@ class GatheredAccount(JMSOrgBaseModel):
def address(self): def address(self):
return self.asset.address 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: class Meta:
verbose_name = _('Gather account automation') verbose_name = _('Gather account automation')
unique_together = [ unique_together = [
@ -31,6 +52,17 @@ class GatheredAccount(JMSOrgBaseModel):
class GatherAccountsAutomation(AccountBaseAutomation): 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): def save(self, *args, **kwargs):
self.type = AutomationTypes.gather_accounts self.type = AutomationTypes.gather_accounts
super().save(*args, **kwargs) super().save(*args, **kwargs)

View File

@ -17,7 +17,8 @@ class GatherAccountAutomationSerializer(BaseAutomationSerializer):
class Meta: class Meta:
model = GatherAccountsAutomation model = GatherAccountsAutomation
read_only_fields = BaseAutomationSerializer.Meta.read_only_fields 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 extra_kwargs = BaseAutomationSerializer.Meta.extra_kwargs