From 4871682dc671276892bb875abf8c95d257f2b054 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@digitalocean.com>
Date: Mon, 18 Dec 2017 16:08:46 -0500
Subject: [PATCH] Allow designating primary IPs assigned to a device's peer VC
 members

---
 netbox/dcim/forms.py  | 22 ++++++++---------
 netbox/dcim/models.py | 56 +++++++++++++++++++++++++------------------
 2 files changed, 43 insertions(+), 35 deletions(-)

diff --git a/netbox/dcim/forms.py b/netbox/dcim/forms.py
index f82787865..b28a1f118 100644
--- a/netbox/dcim/forms.py
+++ b/netbox/dcim/forms.py
@@ -773,26 +773,24 @@ class DeviceForm(BootstrapMixin, TenancyForm, CustomFieldForm):
             # 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.vc_interfaces.values('pk')
+
                 # Collect interface IPs
                 interface_ips = IPAddress.objects.select_related('interface').filter(
-                    family=family, interface__device=self.instance
+                    family=family, interface_id__in=interface_ids
                 )
                 if interface_ips:
-                    ip_choices.append(
-                        ('Interface IPs', [
-                            (ip.id, '{} ({})'.format(ip.address, ip.interface)) for ip in interface_ips
-                        ])
-                    )
+                    ip_list = [(ip.id, '{} ({})'.format(ip.address, ip.interface)) for ip in interface_ips]
+                    ip_choices.append(('Interface IPs', ip_list))
                 # Collect NAT IPs
                 nat_ips = IPAddress.objects.select_related('nat_inside').filter(
-                    family=family, nat_inside__interface__device=self.instance
+                    family=family, nat_inside__interface__in=interface_ids
                 )
                 if nat_ips:
-                    ip_choices.append(
-                        ('NAT IPs', [
-                            (ip.id, '{} ({})'.format(ip.address, ip.nat_inside.address)) for ip in nat_ips
-                        ])
-                    )
+                    ip_list = [(ip.id, '{} ({})'.format(ip.address, ip.nat_inside.address)) for ip in nat_ips]
+                    ip_choices.append(('NAT IPs', ip_list))
                 self.fields['primary_ip{}'.format(family)].choices = ip_choices
 
             # If editing an existing device, exclude it from the list of occupied rack units. This ensures that a device
diff --git a/netbox/dcim/models.py b/netbox/dcim/models.py
index 23b9c0acd..12fd2fe83 100644
--- a/netbox/dcim/models.py
+++ b/netbox/dcim/models.py
@@ -923,29 +923,28 @@ class Device(CreatedUpdatedModel, CustomFieldModel):
             except DeviceType.DoesNotExist:
                 pass
 
-        # Validate primary IPv4 address
-        if self.primary_ip4 and (
-            self.primary_ip4.interface is None or
-            self.primary_ip4.interface.device != self
-        ) and (
-            self.primary_ip4.nat_inside.interface is None or
-            self.primary_ip4.nat_inside.interface.device != self
-        ):
-            raise ValidationError({
-                'primary_ip4': "The specified IP address ({}) is not assigned to this device.".format(self.primary_ip4),
-            })
-
-        # Validate primary IPv6 address
-        if self.primary_ip6 and (
-            self.primary_ip6.interface is None or
-            self.primary_ip6.interface.device != self
-        ) and (
-            self.primary_ip6.nat_inside.interface is None or
-            self.primary_ip6.nat_inside.interface.device != self
-        ):
-            raise ValidationError({
-                'primary_ip6': "The specified IP address ({}) is not assigned to this device.".format(self.primary_ip6),
-            })
+        # Validate primary IP addresses
+        vc_interfaces = self.vc_interfaces.all()
+        if self.primary_ip4:
+            if self.primary_ip4.interface in vc_interfaces:
+                pass
+            elif self.primary_ip4.nat_inside is not None and self.primary_ip4.nat_inside.interface in vc_interfaces:
+                pass
+            else:
+                raise ValidationError({
+                    'primary_ip4': "The specified IP address ({}) is not assigned to this device.".format(
+                        self.primary_ip4),
+                })
+        if self.primary_ip6:
+            if self.primary_ip6.interface in vc_interfaces:
+                pass
+            elif self.primary_ip6.nat_inside is not None and self.primary_ip6.nat_inside.interface in vc_interfaces:
+                pass
+            else:
+                raise ValidationError({
+                    'primary_ip6': "The specified IP address ({}) is not assigned to this device.".format(
+                        self.primary_ip6),
+                })
 
         # A Device can only be assigned to a Cluster in the same Site (or no Site)
         if self.cluster and self.cluster.site is not None and self.cluster.site != self.site:
@@ -1042,6 +1041,17 @@ class Device(CreatedUpdatedModel, CustomFieldModel):
         except VCMembership.DoesNotExist:
             return None
 
+    @property
+    def vc_interfaces(self):
+        """
+        Return a QuerySet matching all Interfaces assigned to this Device or, if this Device is a VC master, to another
+        Device belonging to the same virtual chassis.
+        """
+        if hasattr(self, 'vc_membership') and self.vc_membership.is_master:
+            return Interface.objects.filter(device__vc_membership__virtual_chassis=self.vc_membership.virtual_chassis)
+        else:
+            return self.interfaces.all()
+
     def get_children(self):
         """
         Return the set of child Devices installed in DeviceBays within this Device.