diff --git a/apps/accounts/serializers/account/account.py b/apps/accounts/serializers/account/account.py index f3650e14b..691961f3e 100644 --- a/apps/accounts/serializers/account/account.py +++ b/apps/accounts/serializers/account/account.py @@ -81,21 +81,28 @@ class AccountCreateUpdateSerializerMixin(serializers.Serializer): @staticmethod def get_template_attr_for_account(template): - # Set initial data from template field_names = [ - 'name', 'username', 'secret', 'push_params', - 'secret_type', 'privileged', 'is_active' + 'name', 'username', + 'secret_type', 'secret', + 'privileged', 'is_active' ] + field_map = { + 'push_params': 'params', + 'auto_push': 'push_now' + } + + field_names.extend(field_map.keys()) + attrs = {} for name in field_names: value = getattr(template, name, None) if value is None: continue - if name == 'push_params': - attrs['params'] = value - else: - attrs[name] = value + + attr_name = field_map.get(name, name) + attrs[attr_name] = value + attrs['secret'] = template.get_secret() return attrs @@ -178,7 +185,8 @@ class AccountCreateUpdateSerializerMixin(serializers.Serializer): params = validated_data.pop('params', None) self.clean_auth_fields(validated_data) instance, stat = self.do_create(validated_data) - self.push_account_if_need(instance, push_now, params, stat) + if instance.source == Source.LOCAL: + self.push_account_if_need(instance, push_now, params, stat) return instance def update(self, instance, validated_data): @@ -280,8 +288,8 @@ class AssetAccountBulkSerializer( fields = [ 'name', 'username', 'secret', 'secret_type', 'passphrase', 'privileged', 'is_active', 'comment', 'template', - 'on_invalid', 'push_now', 'assets', 'su_from_username', - 'source', 'source_id', + 'on_invalid', 'push_now', 'params', 'assets', + 'su_from_username', 'source', 'source_id', ] extra_kwargs = { 'name': {'required': False}, @@ -419,16 +427,23 @@ class AssetAccountBulkSerializer( return results @staticmethod - def push_accounts_if_need(results, push_now): + def push_accounts_if_need(results, push_now, params): if not push_now: return - accounts = [str(v['instance']) for v in results if v.get('instance')] - push_accounts_to_assets_task.delay(accounts) + + account_ids = [v['instance'] for v in results if v.get('instance')] + accounts = Account.objects.filter(id__in=account_ids, source=Source.LOCAL) + if not accounts.exists(): + return + + account_ids = [str(_id) for _id in accounts.values_list('id', flat=True)] + push_accounts_to_assets_task.delay(account_ids, params) def create(self, validated_data): + params = validated_data.pop('params', None) push_now = validated_data.pop('push_now', False) results = self.perform_bulk_create(validated_data) - self.push_accounts_if_need(results, push_now) + self.push_accounts_if_need(results, push_now, params) for res in results: res['asset'] = str(res['asset']) return results diff --git a/apps/accounts/signal_handlers.py b/apps/accounts/signal_handlers.py index 3ae1fba5c..ad07c9bb3 100644 --- a/apps/accounts/signal_handlers.py +++ b/apps/accounts/signal_handlers.py @@ -6,6 +6,7 @@ from django.dispatch import receiver from django.utils.translation import gettext_noop from accounts.backends import vault_client +from accounts.const import Source from audits.const import ActivityChoices from audits.signal_handlers import create_activities from common.decorators import merge_delay_run @@ -32,7 +33,7 @@ def push_accounts_if_need(accounts=()): template_accounts = defaultdict(list) for ac in accounts: # 再强调一次吧 - if ac.source != 'template': + if ac.source != Source.TEMPLATE: continue template_accounts[ac.source_id].append(ac) @@ -61,7 +62,7 @@ def create_accounts_activities(account, action='create'): @receiver(post_save, sender=Account) def on_account_create_by_template(sender, instance, created=False, **kwargs): - if not created or instance.source != 'template': + if not created or instance.source != Source.TEMPLATE: return push_accounts_if_need.delay(accounts=(instance,)) create_accounts_activities(instance, action='create') diff --git a/apps/perms/serializers/permission.py b/apps/perms/serializers/permission.py index e4bf33364..38a32a29d 100644 --- a/apps/perms/serializers/permission.py +++ b/apps/perms/serializers/permission.py @@ -6,7 +6,6 @@ from rest_framework import serializers from accounts.const import Source from accounts.models import AccountTemplate, Account -from accounts.tasks import push_accounts_to_assets_task from assets.models import Asset, Node from common.serializers import ResourceLabelsMixin from common.serializers.fields import BitChoicesField, ObjectRelatedField @@ -84,44 +83,35 @@ class AssetPermissionSerializer(ResourceLabelsMixin, BulkOrgResourceModelSeriali return Asset.objects.filter(id__in=asset_ids) def create_accounts(self, assets): - need_create_accounts = [] + account_objs = [] account_attribute = [ 'name', 'username', 'secret_type', 'secret', 'privileged', 'is_active', 'org_id' ] for asset in assets: - asset_exist_accounts = Account.objects.none() - asset_exist_account_names = asset.accounts.values_list('name', flat=True) + asset_exist_account_names = set(asset.accounts.values_list('name', flat=True)) + + asset_exist_accounts = asset.accounts.values('username', 'secret_type') + username_secret_type_set = {(acc['username'], acc['secret_type']) for acc in asset_exist_accounts} for template in self.template_accounts: - asset_exist_accounts |= asset.accounts.filter( - username=template.username, - secret_type=template.secret_type, - ) - username_secret_type_dict = asset_exist_accounts.values('username', 'secret_type') - for template in self.template_accounts: - condition = { - 'username': template.username, - 'secret_type': template.secret_type - } - if condition in username_secret_type_dict or \ - template.name in asset_exist_account_names: + condition = (template.username, template.secret_type) + if condition in username_secret_type_set or template.name in asset_exist_account_names: continue + account_data = {key: getattr(template, key) for key in account_attribute} account_data['su_from'] = template.get_su_from_account(asset) account_data['source'] = Source.TEMPLATE account_data['source_id'] = str(template.id) - need_create_accounts.append(Account(**{'asset_id': asset.id, **account_data})) - return Account.objects.bulk_create(need_create_accounts) + account_objs.append(Account(asset=asset, **account_data)) - def create_and_push_account(self, nodes, assets): + if account_objs: + Account.objects.bulk_create(account_objs) + + def create_account_through_template(self, nodes, assets): if not self.template_accounts: return assets = self.get_all_assets(nodes, assets) - accounts = self.create_accounts(assets) - account_ids = [str(account.id) for account in accounts] - slice_count = 20 - for i in range(0, len(account_ids), slice_count): - push_accounts_to_assets_task.delay(account_ids[i:i + slice_count]) + self.create_accounts(assets) def validate_accounts(self, usernames): template_ids = [] @@ -169,7 +159,7 @@ class AssetPermissionSerializer(ResourceLabelsMixin, BulkOrgResourceModelSeriali instance.nodes.add(*nodes_to_set) def validate(self, attrs): - self.create_and_push_account( + self.create_account_through_template( attrs.get("nodes", []), attrs.get("assets", []) )