From 2ec7ac1ea3b093f761d52e44f627ae1fc037e387 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@digitalocean.com>
Date: Mon, 15 Apr 2019 21:29:02 -0400
Subject: [PATCH] Custom link cleanup

---
 netbox/extras/admin.py                        | 15 +++++++++++
 netbox/extras/constants.py                    | 25 ++++++++++++++++++-
 netbox/extras/migrations/0022_custom_links.py |  2 +-
 netbox/extras/models.py                       | 22 ++++++++++++----
 4 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/netbox/extras/admin.py b/netbox/extras/admin.py
index db193a150..89398ee2b 100644
--- a/netbox/extras/admin.py
+++ b/netbox/extras/admin.py
@@ -81,9 +81,24 @@ class CustomFieldAdmin(admin.ModelAdmin):
 # Custom links
 #
 
+class CustomLinkForm(forms.ModelForm):
+
+    class Meta:
+        model = CustomLink
+        exclude = []
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Format ContentType choices
+        order_content_types(self.fields['content_type'])
+        self.fields['content_type'].choices.insert(0, ('', '---------'))
+
+
 @admin.register(CustomLink, site=admin_site)
 class CustomLinkAdmin(admin.ModelAdmin):
     list_display = ['name', 'content_type', 'group_name', 'weight']
+    form = CustomLinkForm
 
 
 #
diff --git a/netbox/extras/constants.py b/netbox/extras/constants.py
index 3710b3b8d..14c25e2f5 100644
--- a/netbox/extras/constants.py
+++ b/netbox/extras/constants.py
@@ -36,6 +36,29 @@ CF_FILTER_CHOICES = (
 )
 
 # Custom links
+CUSTOM_LINK_MODELS = [
+    'circuits.circuit',
+    'circuits.provider',
+    'dcim.cable',
+    'dcim.device',
+    'dcim.devicetype',
+    'dcim.powerpanel',
+    'dcim.powerfeed',
+    'dcim.rack',
+    'dcim.site',
+    'dcim.region',
+    'ipam.aggregate',
+    'ipam.ipaddress',
+    'ipam.prefix',
+    'ipam.service',
+    'ipam.vlan',
+    'ipam.vrf',
+    'secrets.secret',
+    'tenancy.tenant',
+    'virtualization.cluster',
+    'virtualization.virtualmachine',
+]
+
 BUTTON_CLASS_DEFAULT = 'default'
 BUTTON_CLASS_PRIMARY = 'primary'
 BUTTON_CLASS_SUCCESS = 'success'
@@ -50,7 +73,7 @@ BUTTON_CLASS_CHOICES = (
     (BUTTON_CLASS_INFO, 'Info (aqua)'),
     (BUTTON_CLASS_WARNING, 'Warning (orange)'),
     (BUTTON_CLASS_DANGER, 'Danger (red)'),
-    (BUTTON_CLASS_LINK, 'None'),
+    (BUTTON_CLASS_LINK, 'None (link)'),
 )
 
 # Graph types
diff --git a/netbox/extras/migrations/0022_custom_links.py b/netbox/extras/migrations/0022_custom_links.py
index 24552a2d0..993dc0f42 100644
--- a/netbox/extras/migrations/0022_custom_links.py
+++ b/netbox/extras/migrations/0022_custom_links.py
@@ -21,7 +21,7 @@ class Migration(migrations.Migration):
                 ('url', models.CharField(max_length=200)),
                 ('weight', models.PositiveSmallIntegerField(default=100)),
                 ('group_name', models.CharField(blank=True, max_length=50)),
-                ('button_class', models.CharField(max_length=30)),
+                ('button_class', models.CharField(default='default', max_length=30)),
                 ('new_window', models.BooleanField()),
                 ('content_type', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='contenttypes.ContentType')),
             ],
diff --git a/netbox/extras/models.py b/netbox/extras/models.py
index 0681b882a..0f151fb39 100644
--- a/netbox/extras/models.py
+++ b/netbox/extras/models.py
@@ -304,6 +304,13 @@ class CustomFieldChoice(models.Model):
 # Custom links
 #
 
+def get_custom_link_models():
+    # TODO: This should match on the app_label as well as the model name to avoid potential duplicate names
+    return {
+        'model__in': [model.split('.')[1] for model in CUSTOM_LINK_MODELS],
+    }
+
+
 class CustomLink(models.Model):
     """
     A custom link to an external representation of a NetBox object. The link text and URL fields accept Jinja2 template
@@ -311,7 +318,8 @@ class CustomLink(models.Model):
     """
     content_type = models.ForeignKey(
         to=ContentType,
-        on_delete=models.CASCADE
+        on_delete=models.CASCADE,
+        limit_choices_to=get_custom_link_models
     )
     name = models.CharField(
         max_length=100,
@@ -319,22 +327,26 @@ class CustomLink(models.Model):
     )
     text = models.CharField(
         max_length=200,
-        help_text="Template code for link text"
+        help_text="Jinja2 template code for link text"
     )
     url = models.CharField(
         max_length=200,
-        help_text="Template code for link URL"
+        verbose_name='URL',
+        help_text="Jinja2 template code for link URL"
     )
     weight = models.PositiveSmallIntegerField(
         default=100
     )
     group_name = models.CharField(
         max_length=50,
-        blank=True
+        blank=True,
+        help_text="Links with the same group will appear as a dropdown menu"
     )
     button_class = models.CharField(
         max_length=30,
-        choices=BUTTON_CLASS_CHOICES
+        choices=BUTTON_CLASS_CHOICES,
+        default=BUTTON_CLASS_DEFAULT,
+        help_text="The class of the first link in a group will be used for the dropdown button"
     )
     new_window = models.BooleanField(
         help_text="Force link to open in a new window"