From 6a63b71e15520c1b41996c30ff99ea511e2ada4f Mon Sep 17 00:00:00 2001
From: Daniel Sheppard <dans@dansheps.com>
Date: Thu, 17 Nov 2022 11:01:41 -0600
Subject: [PATCH] Further constrain IP to VDC

---
 netbox/dcim/forms/model_forms.py | 47 +++++++++++++++++++++-----------
 netbox/dcim/models/devices.py    |  5 ++--
 netbox/ipam/models/ip.py         | 11 ++++++--
 3 files changed, 42 insertions(+), 21 deletions(-)

diff --git a/netbox/dcim/forms/model_forms.py b/netbox/dcim/forms/model_forms.py
index 76faa54ba..5758d8caa 100644
--- a/netbox/dcim/forms/model_forms.py
+++ b/netbox/dcim/forms/model_forms.py
@@ -1705,22 +1705,6 @@ class VirtualDeviceContextForm(TenancyForm, NetBoxModelForm):
             'rack_id': '$rack',
         }
     )
-    primary_ip4 = DynamicModelChoiceField(
-        queryset=IPAddress.objects.all(),
-        required=False,
-        query_params={
-            'device_id': '$device',
-            'family': 4,
-        }
-    )
-    primary_ip6 = DynamicModelChoiceField(
-        queryset=IPAddress.objects.all(),
-        required=False,
-        query_params={
-            'device_id': '$device',
-            'family': 6,
-        }
-    )
 
     fieldsets = (
         ('Assigned Device', ('region', 'site_group', 'site', 'location', 'rack', 'device')),
@@ -1736,4 +1720,35 @@ class VirtualDeviceContextForm(TenancyForm, NetBoxModelForm):
         ]
         widgets = {
             'status': StaticSelect(),
+            'primary_ip4': StaticSelect(),
+            'primary_ip6': StaticSelect(),
         }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if self.instance.pk:
+
+            # Compile list of choices for primary IPv4 and IPv6 addresses
+            for family in [4, 6]:
+                ip_choices = [(None, '---------')]
+
+                # Gather PKs of all interfaces belonging to this Device or a peer VirtualChassis member
+                interface_ids = self.instance.interfaces.values_list('pk', flat=True)
+
+                # Collect interface IPs
+                interface_ips = IPAddress.objects.filter(
+                    address__family=family,
+                    assigned_object_type=ContentType.objects.get_for_model(Interface),
+                    assigned_object_id__in=interface_ids
+                ).prefetch_related('assigned_object')
+                if interface_ips:
+                    ip_list = [(ip.id, f'{ip.address} ({ip.assigned_object})') for ip in interface_ips]
+                    ip_choices.append(('Interface IPs', ip_list))
+                self.fields['primary_ip{}'.format(family)].choices = ip_choices
+        else:
+            # An object that doesn't exist yet can't have any IPs assigned to it
+            self.fields['primary_ip4'].choices = []
+            self.fields['primary_ip4'].widget.attrs['readonly'] = True
+            self.fields['primary_ip6'].choices = []
+            self.fields['primary_ip6'].widget.attrs['readonly'] = True
diff --git a/netbox/dcim/models/devices.py b/netbox/dcim/models/devices.py
index 1e6581f8d..c68401ddb 100644
--- a/netbox/dcim/models/devices.py
+++ b/netbox/dcim/models/devices.py
@@ -1178,13 +1178,12 @@ class VirtualDeviceContext(PrimaryModel):
 
     def clean(self):
         super().clean()
-        vc_interfaces = self.device.vc_interfaces(if_master=False)
         if self.primary_ip4:
             if self.primary_ip4.family != 4:
                 raise ValidationError({
                     'primary_ip4': f"{self.primary_ip4} is not an IPv4 address."
                 })
-            if self.primary_ip4.assigned_object not in vc_interfaces:
+            if self.primary_ip4.assigned_object not in self.interfaces.all():
                 raise ValidationError({
                     'primary_ip4': f"The specified IP address ({self.primary_ip4}) is not assigned to this device."
                 })
@@ -1193,7 +1192,7 @@ class VirtualDeviceContext(PrimaryModel):
                 raise ValidationError({
                     'primary_ip6': f"{self.primary_ip6} is not an IPv6 address."
                 })
-            if self.primary_ip6.assigned_object not in vc_interfaces:
+            if self.primary_ip6.assigned_object not in self.interfaces.all():
                 raise ValidationError({
                     'primary_ip6': f"The specified IP address ({self.primary_ip6}) is not assigned to this device."
                 })
diff --git a/netbox/ipam/models/ip.py b/netbox/ipam/models/ip.py
index d5efcdc7f..2f1185e62 100644
--- a/netbox/ipam/models/ip.py
+++ b/netbox/ipam/models/ip.py
@@ -866,9 +866,16 @@ class IPAddress(PrimaryModel):
 
         # Check for primary IP assignment that doesn't match the assigned device/VM
         if self.pk:
-            for cls, attr in ((Device, 'device'), (VirtualMachine, 'virtual_machine'), (VirtualDeviceContext, 'device')):
+            for cls, attr in ((Device, 'device'), (VirtualMachine, 'virtual_machine'), (VirtualDeviceContext, 'vdcs')):
                 parent = cls.objects.filter(Q(primary_ip4=self) | Q(primary_ip6=self)).first()
-                if parent and getattr(self.assigned_object, attr, None) != parent:
+                if parent and hasattr(self.assigned_object, attr) and \
+                    hasattr(getattr(self.assigned_object, attr), 'all') and getattr(self.assigned_object, attr).all() \
+                        != parent:
+                    raise ValidationError({
+                        'interface': f"IP address is primary for {cls._meta.model_name} {parent} but "
+                                     f"not assigned to it!"
+                    })
+                elif parent and getattr(self.assigned_object, attr, None) != parent:
                     # Check for a NAT relationship
                     if not self.nat_inside or getattr(self.nat_inside.assigned_object, attr, None) != parent:
                         raise ValidationError({