From d861d8bfb85871ddde258e78b70e45d1f03cd58e Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@digitalocean.com>
Date: Tue, 2 May 2017 16:46:23 -0400
Subject: [PATCH] Fixes #1118: Allow designating an IP as primary for a device
 while editing the IP

---
 netbox/ipam/forms.py                      | 49 ++++++++++++++++++++++-
 netbox/templates/ipam/ipaddress_edit.html |  1 +
 netbox/utilities/views.py                 |  9 +----
 3 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/netbox/ipam/forms.py b/netbox/ipam/forms.py
index d9c2e30c3..3c47c27df 100644
--- a/netbox/ipam/forms.py
+++ b/netbox/ipam/forms.py
@@ -344,10 +344,11 @@ class IPAddressForm(BootstrapMixin, ReturnURLForm, CustomFieldForm):
             query_key='q', query_url='ipam-api:ipaddress_list', field_to_update='nat_inside', obj_label='address'
         )
     )
+    primary_for_device = forms.BooleanField(required=False, label='Make this the primary IP for the device')
 
     class Meta:
         model = IPAddress
-        fields = ['address', 'vrf', 'tenant', 'status', 'interface', 'nat_inside', 'description']
+        fields = ['address', 'vrf', 'tenant', 'status', 'description', 'interface', 'primary_for_device', 'nat_inside']
         widgets = {
             'interface': APISelect(api_url='/api/dcim/devices/{{interface_device}}/interfaces/'),
             'nat_inside': APISelect(api_url='/api/ipam/ip-addresses/?device_id={{nat_device}}', display_field='address')
@@ -388,6 +389,15 @@ class IPAddressForm(BootstrapMixin, ReturnURLForm, CustomFieldForm):
         else:
             self.fields['interface'].choices = []
 
+        # Initialize primary_for_device if IP address is already assigned
+        if self.instance.interface is not None:
+            device = self.instance.interface.device
+            if (
+                self.instance.address.version == 4 and device.primary_ip4 == self.instance or
+                self.instance.address.version == 6 and device.primary_ip6 == self.instance
+            ):
+                self.initial['primary_for_device'] = True
+
         if self.instance.nat_inside:
             nat_inside = self.instance.nat_inside
             # If the IP is assigned to an interface, populate site/device fields accordingly
@@ -420,6 +430,43 @@ class IPAddressForm(BootstrapMixin, ReturnURLForm, CustomFieldForm):
             else:
                 self.fields['nat_inside'].choices = []
 
+    def clean(self):
+        super(IPAddressForm, self).clean()
+
+        # Primary IP assignment is only available if an interface has been assigned.
+        if self.cleaned_data.get('primary_for_device') and not self.cleaned_data.get('interface'):
+            self.add_error(
+                'primary_for_device', "Only IP addresses assigned to an interface can be designated as primary IPs."
+            )
+
+    def save(self, *args, **kwargs):
+
+        ipaddress = super(IPAddressForm, self).save(*args, **kwargs)
+
+        # Assign this IPAddress as the primary for the associated Device.
+        if self.cleaned_data['primary_for_device']:
+            device = self.cleaned_data['interface'].device
+            if ipaddress.address.version == 4:
+                device.primary_ip4 = ipaddress
+            else:
+                device.primary_ip6 = ipaddress
+            device.save()
+
+        # Clear assignment as primary for device if set.
+        else:
+            try:
+                if ipaddress.address.version == 4:
+                    device = ipaddress.primary_ip4_for
+                    device.primary_ip4 = None
+                else:
+                    device = ipaddress.primary_ip6_for
+                    device.primary_ip6 = None
+                device.save()
+            except Device.DoesNotExist:
+                pass
+
+        return ipaddress
+
 
 class IPAddressBulkAddForm(BootstrapMixin, CustomFieldForm):
     address_pattern = ExpandableIPAddressField(label='Address Pattern')
diff --git a/netbox/templates/ipam/ipaddress_edit.html b/netbox/templates/ipam/ipaddress_edit.html
index 82523eba0..531a91990 100644
--- a/netbox/templates/ipam/ipaddress_edit.html
+++ b/netbox/templates/ipam/ipaddress_edit.html
@@ -28,6 +28,7 @@
             {% render_field form.interface_rack %}
             {% render_field form.interface_device %}
             {% render_field form.interface %}
+            {% render_field form.primary_for_device %}
         </div>
     </div>
     <div class="panel panel-default">
diff --git a/netbox/utilities/views.py b/netbox/utilities/views.py
index 9ed35aaa0..80539a441 100644
--- a/netbox/utilities/views.py
+++ b/netbox/utilities/views.py
@@ -17,7 +17,6 @@ from django.utils.http import is_safe_url
 from django.utils.safestring import mark_safe
 from django.views.generic import View
 
-from extras.forms import CustomFieldForm
 from extras.models import CustomField, CustomFieldValue, ExportTemplate, UserAction
 
 from .error_handlers import handle_protectederror
@@ -195,12 +194,8 @@ class ObjectEditView(GetReturnURLMixin, View):
         form = self.form_class(request.POST, instance=obj)
 
         if form.is_valid():
-            obj = form.save(commit=False)
-            obj_created = not obj.pk
-            obj.save()
-            form.save_m2m()
-            if isinstance(form, CustomFieldForm):
-                form.save_custom_fields()
+            obj_created = not form.instance.pk
+            obj = form.save()
 
             msg = u'Created ' if obj_created else u'Modified '
             msg += self.model._meta.verbose_name