# -*- coding: utf-8 -*- # from django.db.models import Q from django.utils.translation import ugettext_lazy as _ from rest_framework import serializers from accounts.models import AccountTemplate, Account from accounts.tasks import push_accounts_to_assets_task from assets.models import Asset, Node from common.serializers.fields import BitChoicesField, ObjectRelatedField from orgs.mixins.serializers import BulkOrgResourceModelSerializer from perms.models import ActionChoices, AssetPermission from users.models import User, UserGroup __all__ = ["AssetPermissionSerializer", "ActionChoicesField"] class ActionChoicesField(BitChoicesField): def __init__(self, **kwargs): super().__init__(choice_cls=ActionChoices, **kwargs) def to_file_representation(self, value): return [v['value'] for v in value] def to_file_internal_value(self, data): return data class AssetPermissionSerializer(BulkOrgResourceModelSerializer): users = ObjectRelatedField(queryset=User.objects, many=True, required=False, label=_('User')) user_groups = ObjectRelatedField( queryset=UserGroup.objects, many=True, required=False, label=_('User group') ) assets = ObjectRelatedField(queryset=Asset.objects, many=True, required=False, label=_('Asset')) nodes = ObjectRelatedField(queryset=Node.objects, many=True, required=False, label=_('Node')) actions = ActionChoicesField(required=False, allow_null=True, label=_("Actions")) is_valid = serializers.BooleanField(read_only=True, label=_("Is valid")) is_expired = serializers.BooleanField(read_only=True, label=_("Is expired")) accounts = serializers.ListField(label=_("Account"), required=False) template_accounts = AccountTemplate.objects.none() class Meta: model = AssetPermission fields_mini = ["id", "name"] fields_generic = [ "accounts", "actions", "created_by", "date_created", "date_start", "date_expired", "is_active", "is_expired", "is_valid", "comment", "from_ticket", ] fields_small = fields_mini + fields_generic fields_m2m = [ "users", "user_groups", "assets", "nodes", ] fields = fields_mini + fields_m2m + fields_generic read_only_fields = ["created_by", "date_created", "from_ticket"] extra_kwargs = { "actions": {"label": _("Actions")}, "is_expired": {"label": _("Is expired")}, "is_valid": {"label": _("Is valid")}, } def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.set_actions_field() def set_actions_field(self): actions = self.fields.get("actions") if not actions: return actions.default = list(actions.choices.keys()) @staticmethod def get_all_assets(nodes, assets): node_asset_ids = Node.get_nodes_all_assets(*nodes).values_list('id', flat=True) direct_asset_ids = [asset.id for asset in assets] asset_ids = set(direct_asset_ids + list(node_asset_ids)) return Asset.objects.filter(id__in=asset_ids) def create_accounts(self, assets): need_create_accounts = [] account_attribute = [ 'name', 'username', 'secret_type', 'secret', 'privileged', 'is_active', 'org_id' ] for asset in assets: asset_exist_accounts = Account.objects.none() 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: continue account_data = {key: getattr(template, key) for key in account_attribute} account_data['name'] = f"{account_data['name']}-{_('Account template')}" need_create_accounts.append(Account(**{'asset_id': asset.id, **account_data})) return Account.objects.bulk_create(need_create_accounts) def create_and_push_account(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]) def validate_accounts(self, usernames: list[str]): template_ids = [] account_usernames = [] for username in usernames: if username.startswith('%'): template_ids.append(username[1:]) else: account_usernames.append(username) self.template_accounts = AccountTemplate.objects.filter(id__in=template_ids) template_usernames = list(self.template_accounts.values_list('username', flat=True)) return list(set(account_usernames + template_usernames)) @classmethod def setup_eager_loading(cls, queryset): """Perform necessary eager loading of data.""" queryset = queryset.prefetch_related( "users", "user_groups", "assets", "nodes", ) return queryset @staticmethod def perform_display_create(instance, **kwargs): # 用户 users_to_set = User.objects.filter( Q(name__in=kwargs.get("users_display")) | Q(username__in=kwargs.get("users_display")) ).distinct() instance.users.add(*users_to_set) # 用户组 user_groups_to_set = UserGroup.objects.filter( name__in=kwargs.get("user_groups_display") ).distinct() instance.user_groups.add(*user_groups_to_set) # 资产 assets_to_set = Asset.objects.filter( Q(address__in=kwargs.get("assets_display")) | Q(name__in=kwargs.get("assets_display")) ).distinct() instance.assets.add(*assets_to_set) # 节点 nodes_to_set = Node.objects.filter( full_value__in=kwargs.get("nodes_display") ).distinct() instance.nodes.add(*nodes_to_set) def validate(self, attrs): self.create_and_push_account( attrs.get("nodes", []), attrs.get("assets", []) ) return super().validate(attrs) def create(self, validated_data): display = { "users_display": validated_data.pop("users_display", ""), "user_groups_display": validated_data.pop("user_groups_display", ""), "assets_display": validated_data.pop("assets_display", ""), "nodes_display": validated_data.pop("nodes_display", ""), } instance = super().create(validated_data) self.perform_display_create(instance, **display) return instance