jumpserver/apps/assets/serializers/domain.py

# -*- coding: utf-8 -*-
#
from rest_framework import serializers
from rest_framework.generics import get_object_or_404
from django.utils.translation import ugettext_lazy as _

from orgs.mixins.serializers import BulkOrgResourceModelSerializer
from common.drf.serializers import SecretReadableMixin
from common.drf.fields import ObjectRelatedField, EncryptedField
from assets.const import SecretType
from ..models import Domain, Asset, Account
from ..serializers import HostSerializer
from .utils import validate_password_for_ansible, validate_ssh_key


class DomainSerializer(BulkOrgResourceModelSerializer):
    asset_count = serializers.SerializerMethodField(label=_('Assets amount'))
    gateway_count = serializers.SerializerMethodField(label=_('Gateways count'))
    assets = ObjectRelatedField(
        many=True, required=False, queryset=Asset.objects, label=_('Asset')
    )

    class Meta:
        model = Domain
        fields_mini = ['id', 'name']
        fields_small = fields_mini + ['comment']
        fields_m2m = ['assets']
        read_only_fields = ['asset_count', 'gateway_count', 'date_created']
        fields = fields_small + fields_m2m + read_only_fields

        extra_kwargs = {
            'assets': {'required': False, 'label': _('Assets')},
        }

    @staticmethod
    def get_asset_count(obj):
        return obj.assets.count()

    @staticmethod
    def get_gateway_count(obj):
        return obj.gateways.count()


class GatewaySerializer(HostSerializer):
    password = EncryptedField(
        label=_('Password'), required=False, allow_blank=True, allow_null=True, max_length=1024,
        validators=[validate_password_for_ansible], write_only=True
    )
    private_key = EncryptedField(
        label=_('SSH private key'), required=False, allow_blank=True, allow_null=True,
        max_length=16384, write_only=True
    )
    passphrase = serializers.CharField(
        label=_('Key password'), allow_blank=True, allow_null=True, required=False, write_only=True,
        max_length=512,
    )
    username = serializers.CharField(
        label=_('Username'), allow_blank=True, max_length=128, required=True,
    )

    class Meta(HostSerializer.Meta):
        fields = HostSerializer.Meta.fields + [
            'username', 'password', 'private_key', 'passphrase'
        ]

    def validate_private_key(self, secret):
        if not secret:
            return
        passphrase = self.initial_data.get('passphrase')
        passphrase = passphrase if passphrase else None
        validate_ssh_key(secret, passphrase)
        return secret

    @staticmethod
    def clean_auth_fields(validated_data):
        username = validated_data.pop('username', None)
        password = validated_data.pop('password', None)
        private_key = validated_data.pop('private_key', None)
        validated_data.pop('passphrase', None)
        return username, password, private_key

    @staticmethod
    def create_accounts(instance, username, password, private_key):
        account_name = f'{instance.name}-{_("Gateway")}'
        account_data = {
            'privileged': True,
            'name': account_name,
            'username': username,
            'asset_id': instance.id,
            'created_by': instance.created_by
        }
        if password:
            Account.objects.create(
                **account_data, secret=password, secret_type=SecretType.PASSWORD
            )
        if private_key:
            Account.objects.create(
                **account_data, secret=private_key, secret_type=SecretType.SSH_KEY
            )

    @staticmethod
    def update_accounts(instance, username, password, private_key):
        accounts = instance.accounts.filter(username=username)
        if password:
            account = get_object_or_404(accounts, SecretType.PASSWORD)
            account.secret = password
            account.save()
        if private_key:
            account = get_object_or_404(accounts, SecretType.SSH_KEY)
            account.secret = private_key
            account.save()

    def create(self, validated_data):
        auth_fields = self.clean_auth_fields(validated_data)
        instance = super().create(validated_data)
        self.create_accounts(instance, *auth_fields)
        return instance

    def update(self, instance, validated_data):
        auth_fields = self.clean_auth_fields(validated_data)
        instance = super().update(instance, validated_data)
        self.update_accounts(instance, *auth_fields)
        return instance


class GatewayWithAuthSerializer(SecretReadableMixin, GatewaySerializer):
    class Meta(GatewaySerializer.Meta):
        extra_kwargs = {
            'password': {'write_only': False},
            'private_key': {"write_only": False},
            'public_key': {"write_only": False},
        }


class DomainWithGatewaySerializer(BulkOrgResourceModelSerializer):
    gateways = GatewayWithAuthSerializer(many=True, read_only=True)

    class Meta:
        model = Domain
        fields = '__all__'