netbox/netbox/vpn/forms/bulk_import.py

from django.core.exceptions import ValidationError
from django.utils.translation import gettext_lazy as _

from dcim.models import Device, Interface
from ipam.models import VLAN, VRF, IPAddress
from netbox.forms import NetBoxModelImportForm, OrganizationalModelImportForm, PrimaryModelImportForm
from tenancy.models import Tenant
from utilities.forms.fields import CSVChoiceField, CSVModelChoiceField, CSVModelMultipleChoiceField
from virtualization.models import VirtualMachine, VMInterface
from vpn.choices import *
from vpn.models import *

__all__ = (
    'IKEPolicyImportForm',
    'IKEProposalImportForm',
    'IPSecPolicyImportForm',
    'IPSecProfileImportForm',
    'IPSecProposalImportForm',
    'L2VPNImportForm',
    'L2VPNTerminationImportForm',
    'TunnelGroupImportForm',
    'TunnelImportForm',
    'TunnelTerminationImportForm',
)


class TunnelGroupImportForm(OrganizationalModelImportForm):

    class Meta:
        model = TunnelGroup
        fields = ('name', 'slug', 'description', 'owner', 'comments', 'tags')


class TunnelImportForm(PrimaryModelImportForm):
    status = CSVChoiceField(
        label=_('Status'),
        choices=TunnelStatusChoices,
        help_text=_('Operational status')
    )
    group = CSVModelChoiceField(
        label=_('Tunnel group'),
        queryset=TunnelGroup.objects.all(),
        required=False,
        to_field_name='name'
    )
    encapsulation = CSVChoiceField(
        label=_('Encapsulation'),
        choices=TunnelEncapsulationChoices,
        help_text=_('Tunnel encapsulation')
    )
    ipsec_profile = CSVModelChoiceField(
        label=_('IPSec profile'),
        queryset=IPSecProfile.objects.all(),
        required=False,
        to_field_name='name'
    )
    tenant = CSVModelChoiceField(
        label=_('Tenant'),
        queryset=Tenant.objects.all(),
        required=False,
        to_field_name='name',
        help_text=_('Assigned tenant')
    )

    class Meta:
        model = Tunnel
        fields = (
            'name', 'status', 'group', 'encapsulation', 'ipsec_profile', 'tenant', 'tunnel_id', 'description',
            'owner', 'comments', 'tags',
        )


class TunnelTerminationImportForm(NetBoxModelImportForm):
    tunnel = CSVModelChoiceField(
        label=_('Tunnel'),
        queryset=Tunnel.objects.all(),
        to_field_name='name'
    )
    role = CSVChoiceField(
        label=_('Role'),
        choices=TunnelTerminationRoleChoices,
        help_text=_('Operational role')
    )
    device = CSVModelChoiceField(
        label=_('Device'),
        queryset=Device.objects.all(),
        required=False,
        to_field_name='name',
        help_text=_('Parent device of assigned interface')
    )
    virtual_machine = CSVModelChoiceField(
        label=_('Virtual machine'),
        queryset=VirtualMachine.objects.all(),
        required=False,
        to_field_name='name',
        help_text=_('Parent VM of assigned interface')
    )
    termination = CSVModelChoiceField(
        label=_('Termination'),
        queryset=Interface.objects.none(),  # Can also refer to VMInterface
        required=False,
        to_field_name='name',
        help_text=_('Device or virtual machine interface')
    )
    outside_ip = CSVModelChoiceField(
        label=_('Outside IP'),
        queryset=IPAddress.objects.all(),
        required=False,
        to_field_name='address'
    )

    class Meta:
        model = TunnelTermination
        fields = (
            'tunnel', 'role', 'outside_ip', 'tags',
        )

    def __init__(self, data=None, *args, **kwargs):
        super().__init__(data, *args, **kwargs)

        if data:

            # Limit termination queryset by assigned device/VM
            if data.get('device'):
                self.fields['termination'].queryset = Interface.objects.filter(
                    **{f"device__{self.fields['device'].to_field_name}": data['device']}
                )
            elif data.get('virtual_machine'):
                self.fields['termination'].queryset = VMInterface.objects.filter(
                    **{f"virtual_machine__{self.fields['virtual_machine'].to_field_name}": data['virtual_machine']}
                )

    def save(self, *args, **kwargs):

        # Assign termination object
        if self.cleaned_data.get('termination'):
            self.instance.termination = self.cleaned_data['termination']

        return super().save(*args, **kwargs)


class IKEProposalImportForm(PrimaryModelImportForm):
    authentication_method = CSVChoiceField(
        label=_('Authentication method'),
        choices=AuthenticationMethodChoices
    )
    encryption_algorithm = CSVChoiceField(
        label=_('Encryption algorithm'),
        choices=EncryptionAlgorithmChoices
    )
    authentication_algorithm = CSVChoiceField(
        label=_('Authentication algorithm'),
        choices=AuthenticationAlgorithmChoices,
        required=False
    )
    group = CSVChoiceField(
        label=_('Group'),
        choices=DHGroupChoices
    )

    class Meta:
        model = IKEProposal
        fields = (
            'name', 'description', 'authentication_method', 'encryption_algorithm', 'authentication_algorithm',
            'group', 'sa_lifetime', 'owner', 'comments', 'tags',
        )


class IKEPolicyImportForm(PrimaryModelImportForm):
    version = CSVChoiceField(
        label=_('Version'),
        choices=IKEVersionChoices
    )
    mode = CSVChoiceField(
        label=_('Mode'),
        choices=IKEModeChoices,
        required=False
    )
    proposals = CSVModelMultipleChoiceField(
        queryset=IKEProposal.objects.all(),
        to_field_name='name',
        help_text=_('IKE proposal(s)'),
    )

    class Meta:
        model = IKEPolicy
        fields = (
            'name', 'description', 'version', 'mode', 'proposals', 'preshared_key', 'owner', 'comments', 'tags',
        )


class IPSecProposalImportForm(PrimaryModelImportForm):
    encryption_algorithm = CSVChoiceField(
        label=_('Encryption algorithm'),
        choices=EncryptionAlgorithmChoices,
        required=False
    )
    authentication_algorithm = CSVChoiceField(
        label=_('Authentication algorithm'),
        choices=AuthenticationAlgorithmChoices,
        required=False
    )

    class Meta:
        model = IPSecProposal
        fields = (
            'name', 'description', 'encryption_algorithm', 'authentication_algorithm', 'sa_lifetime_seconds',
            'sa_lifetime_data', 'owner', 'comments', 'tags',
        )


class IPSecPolicyImportForm(PrimaryModelImportForm):
    pfs_group = CSVChoiceField(
        label=_('Diffie-Hellman group for Perfect Forward Secrecy'),
        choices=DHGroupChoices,
        required=False
    )
    proposals = CSVModelMultipleChoiceField(
        queryset=IPSecProposal.objects.all(),
        to_field_name='name',
        help_text=_('IPSec proposal(s)'),
    )

    class Meta:
        model = IPSecPolicy
        fields = (
            'name', 'description', 'proposals', 'pfs_group', 'owner', 'comments', 'tags',
        )


class IPSecProfileImportForm(PrimaryModelImportForm):
    mode = CSVChoiceField(
        label=_('Mode'),
        choices=IPSecModeChoices,
        help_text=_('IPSec protocol')
    )
    ike_policy = CSVModelChoiceField(
        label=_('IKE policy'),
        queryset=IKEPolicy.objects.all(),
        to_field_name='name'
    )
    ipsec_policy = CSVModelChoiceField(
        label=_('IPSec policy'),
        queryset=IPSecPolicy.objects.all(),
        to_field_name='name'
    )

    class Meta:
        model = IPSecProfile
        fields = (
            'name', 'mode', 'ike_policy', 'ipsec_policy', 'description', 'owner', 'comments', 'tags',
        )


class L2VPNImportForm(PrimaryModelImportForm):
    tenant = CSVModelChoiceField(
        label=_('Tenant'),
        queryset=Tenant.objects.all(),
        required=False,
        to_field_name='name',
    )
    status = CSVChoiceField(
        label=_('Status'),
        choices=L2VPNStatusChoices,
        help_text=_('Operational status')
    )
    type = CSVChoiceField(
        label=_('Type'),
        choices=L2VPNTypeChoices,
        help_text=_('L2VPN type')
    )

    class Meta:
        model = L2VPN
        fields = (
            'identifier', 'name', 'slug', 'tenant', 'type', 'description', 'owner', 'comments', 'tags',
        )


class L2VPNTerminationImportForm(NetBoxModelImportForm):
    l2vpn = CSVModelChoiceField(
        queryset=L2VPN.objects.all(),
        required=True,
        to_field_name='name',
        label=_('L2VPN'),
    )
    device = CSVModelChoiceField(
        label=_('Device'),
        queryset=Device.objects.all(),
        required=False,
        to_field_name='name',
        help_text=_('Parent device (for interface or IP address)')
    )
    virtual_machine = CSVModelChoiceField(
        label=_('Virtual machine'),
        queryset=VirtualMachine.objects.all(),
        required=False,
        to_field_name='name',
        help_text=_('Parent virtual machine (for interface or IP address)')
    )
    interface = CSVModelChoiceField(
        label=_('Interface'),
        queryset=Interface.objects.none(),  # Can also refer to VMInterface
        required=False,
        to_field_name='name',
        help_text=_('Assigned interface (device or VM)')
    )
    vlan = CSVModelChoiceField(
        label=_('VLAN'),
        queryset=VLAN.objects.all(),
        required=False,
        to_field_name='name',
        help_text=_('Assigned VLAN')
    )
    ipaddress = CSVModelChoiceField(
        label=_('IP address'),
        queryset=IPAddress.objects.all(),
        required=False,
        to_field_name='address',
        help_text=_('Assigned IP address, with mask (e.g. the far-end PE of a pseudowire)')
    )
    ipaddress_vrf = CSVModelChoiceField(
        label=_('IP address VRF'),
        queryset=VRF.objects.all(),
        required=False,
        to_field_name='name',
        help_text=_('VRF of the assigned IP address')
    )

    class Meta:
        model = L2VPNTermination
        fields = ('l2vpn', 'device', 'virtual_machine', 'interface', 'vlan', 'ipaddress', 'ipaddress_vrf', 'tags')

    def __init__(self, data=None, *args, **kwargs):
        super().__init__(data, *args, **kwargs)

        if data:

            # Limit interface and IP address querysets by device or VM
            if data.get('device'):
                self.fields['interface'].queryset = Interface.objects.filter(
                    **{f"device__{self.fields['device'].to_field_name}": data['device']}
                )
                self.fields['ipaddress'].queryset = IPAddress.objects.filter(
                    **{f"interface__device__{self.fields['device'].to_field_name}": data['device']}
                )
            elif data.get('virtual_machine'):
                self.fields['interface'].queryset = VMInterface.objects.filter(
                    **{f"virtual_machine__{self.fields['virtual_machine'].to_field_name}": data['virtual_machine']}
                )
                self.fields['ipaddress'].queryset = IPAddress.objects.filter(
                    **{
                        f"vminterface__virtual_machine__{self.fields['virtual_machine'].to_field_name}":
                            data['virtual_machine']
                    }
                )

            # Limit IP address queryset by VRF
            if data.get('ipaddress_vrf'):
                self.fields['ipaddress'].queryset = self.fields['ipaddress'].queryset.filter(
                    **{f"vrf__{self.fields['ipaddress_vrf'].to_field_name}": data['ipaddress_vrf']}
                )

    def clean(self):
        super().clean()

        if self.cleaned_data.get('device') and self.cleaned_data.get('virtual_machine'):
            raise ValidationError(_('Cannot import device and VM interface terminations simultaneously.'))
        assigned_objects = [
            obj for obj in (
                self.cleaned_data.get('interface'),
                self.cleaned_data.get('vlan'),
                self.cleaned_data.get('ipaddress'),
            ) if obj
        ]
        if not self.instance.pk and not assigned_objects:
            raise ValidationError(_('Each termination must specify an interface, a VLAN, or an IP address.'))
        if len(assigned_objects) > 1:
            raise ValidationError(_('Cannot assign more than one of an interface, a VLAN, and an IP address.'))

        # if this is an update we might not have an assigned object in the form data
        if assigned_objects:
            self.instance.assigned_object = assigned_objects[0]