Generic Relation for Secrets

netbox/dcim/api/serializers.py

@@ -17,6 +17,7 @@ from dcim.models import (
 )
 from extras.api.customfields import CustomFieldModelSerializer
 from ipam.models import IPAddress, VLAN
+from secrets.models import Secret
 from tenancy.api.serializers import NestedTenantSerializer
 from users.api.serializers import NestedUserSerializer
 from utilities.api import (

netbox/dcim/models.py

@@ -1287,6 +1287,9 @@ class Device(ChangeLoggedModel, ConfigContextModel, CustomFieldModel):
     images = GenericRelation(
         to='extras.ImageAttachment'
     )
+    secrets = GenericRelation(
+        to='secret.Secret'
+    )
 
     objects = DeviceManager()
     tags = TaggableManager()

netbox/secrets/api/serializers.py

@@ -5,6 +5,7 @@ from rest_framework.validators import UniqueTogetherValidator
 from taggit_serializer.serializers import TaggitSerializer, TagListSerializerField
 
 from dcim.api.serializers import NestedDeviceSerializer
+from dcim.models import Device
 from extras.api.customfields import CustomFieldModelSerializer
 from secrets.models import Secret, SecretRole
 from utilities.api import ValidatedModelSerializer, WritableNestedSerializer
@@ -33,8 +34,19 @@ class NestedSecretRoleSerializer(WritableNestedSerializer):
 # Secrets
 #
 
+class SecretObjectRelatedField(serializers.RelatedField):
+    """
+        GenericRelation object can be extended with different Object types
+    """
+    def to_representation(self, value):
+        if isinstance(value, Device):
+            object_serial = NestedDeviceSerializer(value, context=self.context)
+            dict = object_serial.data.copy()
+            dict.update({'class': 'device'})
+            return dict
+
 class SecretSerializer(TaggitSerializer, CustomFieldModelSerializer):
-    device = NestedDeviceSerializer()
+    object = SecretObjectRelatedField(queryset=Secret.objects.all(), required=False)
     role = NestedSecretRoleSerializer()
     plaintext = serializers.CharField()
     tags = TagListSerializerField(required=False)
@@ -42,12 +54,11 @@ class SecretSerializer(TaggitSerializer, CustomFieldModelSerializer):
     class Meta:
         model = Secret
         fields = [
-            'id', 'device', 'role', 'name', 'plaintext', 'hash', 'tags', 'custom_fields', 'created', 'last_updated',
+            'id', 'object', 'role', 'name', 'plaintext', 'hash', 'tags', 'custom_fields', 'created', 'last_updated',
         ]
         validators = []
 
     def validate(self, data):
-
         # Encrypt plaintext data using the master key provided from the view context
         if data.get('plaintext'):
             s = Secret(plaintext=data['plaintext'])
@@ -57,7 +68,7 @@ class SecretSerializer(TaggitSerializer, CustomFieldModelSerializer):
 
         # Validate uniqueness of name if one has been provided.
         if data.get('name'):
-            validator = UniqueTogetherValidator(queryset=Secret.objects.all(), fields=('device', 'role', 'name'))
+            validator = UniqueTogetherValidator(queryset=Secret.objects.all(), fields=('object', 'role', 'name'))
             validator.set_context(self)
             validator(data)
 

netbox/secrets/api/views.py

@@ -46,7 +46,7 @@ class SecretRoleViewSet(ModelViewSet):
 
 class SecretViewSet(ModelViewSet):
     queryset = Secret.objects.select_related(
-        'device__primary_ip4', 'device__primary_ip6', 'role',
+        'role',
     ).prefetch_related(
         'role__users', 'role__groups',
     )

netbox/secrets/filters.py

@@ -34,7 +34,7 @@ class SecretFilter(CustomFieldFilterSet, django_filters.FilterSet):
     )
     device_id = django_filters.ModelMultipleChoiceFilter(
         queryset=Device.objects.all(),
-        label='Device (ID)',
+        label='Object (ID)',
     )
     device = django_filters.ModelMultipleChoiceFilter(
         name='device__name',

netbox/secrets/migrations/0006_foregin_key_generic.py

@@ -0,0 +1,20 @@
+# Generated by Django 2.0.8 on 2018-08-10 15:28
+ from django.db import migrations, models
+import django.db.models.deletion
+ class Migration(migrations.Migration):
+     dependencies = [
+        ('contenttypes', '0002_remove_content_type_name'),
+        ('secrets', '0005_change_logging'),
+    ]
+     operations = [
+        migrations.AddField(
+            model_name='secret',
+            name='content_type',
+            field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, to='contenttypes.ContentType'),
+        ),
+        migrations.AddField(
+            model_name='secret',
+            name='object_id',
+            field=models.PositiveIntegerField(blank=True, null=True),
+        ),
+    ]

netbox/secrets/migrations/0007_populate_generic_foregin_key.py

@@ -0,0 +1,22 @@
+# Generated by Django 2.0.8 on 2018-08-10 15:29
+ from django.db import migrations
+ def migrate_device_secret(apps, schema_editor):
+    """
+        Move data from device foreign key into GenericForeignKey
+    """
+    Secret = apps.get_model('secrets', 'Secret')
+    if Secret.objects.count() > 0:
+        ContentType = apps.get_model('contenttypes', 'ContentType')
+         secret_device_content_type = ContentType.objects.get(app_label='dcim', model='device')
+         for secret in Secret.objects.all():
+            secret.content_type = secret_device_content_type
+            secret.object_id = secret.device.pk
+            secret.save()
+ class Migration(migrations.Migration):
+     dependencies = [
+        ('dcim', '0058_relax_rack_naming_constraints'),
+        ('secrets', '0006_foreign_key_generic'),
+    ]
+     operations = [
+        migrations.RunPython(migrate_device_secret)
+    ]

netbox/secrets/migrations/0008_create_related_object.py

@@ -0,0 +1,32 @@
+# Generated by Django 2.0.8 on 2018-08-10 16:37
+ from django.db import migrations, models
+import django.db.models.deletion
+ class Migration(migrations.Migration):
+     dependencies = [
+        ('contenttypes', '0002_remove_content_type_name'),
+        ('secrets', '0007_populate_generic_foreign_key'),
+    ]
+     operations = [
+        migrations.AlterModelOptions(
+            name='secret',
+            options={'ordering': ['content_type', 'object_id', 'role', 'name']},
+        ),
+        migrations.AlterField(
+            model_name='secret',
+            name='content_type',
+            field=models.ForeignKey(default=33, on_delete=django.db.models.deletion.CASCADE, to='contenttypes.ContentType'),
+        ),
+        migrations.AlterField(
+            model_name='secret',
+            name='object_id',
+            field=models.PositiveIntegerField(default=1),
+        ),
+        migrations.AlterUniqueTogether(
+            name='secret',
+            unique_together={('content_type', 'object_id', 'role', 'name')},
+        ),
+        migrations.RemoveField(
+            model_name='secret',
+            name='device',
+        ),
+     ]

netbox/secrets/models.py

@@ -8,7 +8,9 @@ from Crypto.Util import strxor
 from django.conf import settings
 from django.contrib.auth.hashers import make_password, check_password
 from django.contrib.auth.models import Group, User
+from django.contrib.contenttypes.fields import GenericForeignKey
 from django.contrib.contenttypes.fields import GenericRelation
+from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ValidationError
 from django.db import models
 from django.urls import reverse
@@ -322,11 +324,6 @@ class Secret(ChangeLoggedModel, CustomFieldModel):
     A Secret can be up to 65,536 bytes (64KB) in length. Each secret string will be padded with random data to a minimum
     of 64 bytes during encryption in order to protect short strings from ciphertext analysis.
     """
-    device = models.ForeignKey(
-        to='dcim.Device',
-        on_delete=models.CASCADE,
-        related_name='secrets'
-    )
     role = models.ForeignKey(
         to='secrets.SecretRole',
         on_delete=models.PROTECT,
@@ -349,39 +346,60 @@ class Secret(ChangeLoggedModel, CustomFieldModel):
         content_type_field='obj_type',
         object_id_field='obj_id'
     )
+    content_type = models.ForeignKey(
+        ContentType,
+        on_delete=models.CASCADE,
+        default=ContentType.objects.get(app_label='dcim', model='device').pk
+    )
+    object_id = models.PositiveIntegerField(default=1)
+    object = GenericForeignKey('content_type', 'object_id')
 
     tags = TaggableManager()
 
     plaintext = None
-    csv_headers = ['device', 'role', 'name', 'plaintext']
+    csv_headers = ['object_id', 'content_type', 'role', 'name', 'plaintext']
 
     class Meta:
-        ordering = ['device', 'role', 'name']
+        ordering = ['content_type', 'object_id', 'role', 'name']
-        unique_together = ['device', 'role', 'name']
+        unique_together = ['content_type', 'object_id', 'role', 'name']
 
     def __init__(self, *args, **kwargs):
         self.plaintext = kwargs.pop('plaintext', None)
         super(Secret, self).__init__(*args, **kwargs)
 
     def __str__(self):
-        if self.role and self.device and self.name:
+        if self.role and self.object and self.name:
-            return '{} for {} ({})'.format(self.role, self.device, self.name)
+            return '{} for {} ({})'.format(self.role, self.object, self.name)
         # Return role and device if no name is set
-        if self.role and self.device:
+        if self.role and self.object:
-            return '{} for {}'.format(self.role, self.device)
+            return '{} for {}'.format(self.role, self.object)
         return 'Secret'
 
+    def get_device_secrets(device):
+        """
+            Returns Secrets given a device
+        """
+        return Secret.objects.filter(
+            content_type=ContentType.objects.get(
+                app_label='dcim', model='device'
+            ),
+            object_id=device.pk
+        )
+
     def get_absolute_url(self):
         return reverse('secrets:secret', args=[self.pk])
 
     def to_csv(self):
         return (
-            self.device,
+            self.object,
             self.role,
             self.name,
             self.plaintext or '',
         )
 
+    def object_label(self):
+        return(str(self.content_type).capitalize())
+
     def _pad(self, s):
         """
         Prepend the length of the plaintext (2B) and pad with garbage to a multiple of 16B (minimum of 64B).

netbox/secrets/tables.py

@@ -39,8 +39,8 @@ class SecretRoleTable(BaseTable):
 
 class SecretTable(BaseTable):
     pk = ToggleColumn()
-    device = tables.LinkColumn()
+    object = tables.LinkColumn()
 
     class Meta(BaseTable.Meta):
         model = Secret
-        fields = ('pk', 'device', 'role', 'name', 'last_updated')
+        fields = ('pk', 'object', 'role', 'name', 'last_updated')

netbox/secrets/tests/test_api.py

@@ -173,17 +173,17 @@ class SecretTest(APITestCase):
         self.secretrole1 = SecretRole.objects.create(name='Test Secret Role 1', slug='test-secret-role-1')
         self.secretrole2 = SecretRole.objects.create(name='Test Secret Role 2', slug='test-secret-role-2')
         self.secret1 = Secret(
-            device=self.device, role=self.secretrole1, name='Test Secret 1', plaintext=self.plaintext['secret1']
+            object=self.device, role=self.secretrole1, name='Test Secret 1', plaintext=self.plaintext['secret1']
         )
         self.secret1.encrypt(self.master_key)
         self.secret1.save()
         self.secret2 = Secret(
-            device=self.device, role=self.secretrole1, name='Test Secret 2', plaintext=self.plaintext['secret2']
+            object=self.device, role=self.secretrole1, name='Test Secret 2', plaintext=self.plaintext['secret2']
         )
         self.secret2.encrypt(self.master_key)
         self.secret2.save()
         self.secret3 = Secret(
-            device=self.device, role=self.secretrole1, name='Test Secret 3', plaintext=self.plaintext['secret3']
+            object=self.device, role=self.secretrole1, name='Test Secret 3', plaintext=self.plaintext['secret3']
         )
         self.secret3.encrypt(self.master_key)
         self.secret3.save()
@@ -205,7 +205,8 @@ class SecretTest(APITestCase):
     def test_create_secret(self):
 
         data = {
-            'device': self.device.pk,
+            'object_id': self.device.pk,
+            'content_type': '33',
             'role': self.secretrole1.pk,
             'name': 'Test Secret 4',
             'plaintext': 'Secret #4 Plaintext',
@@ -226,19 +227,22 @@ class SecretTest(APITestCase):
 
         data = [
             {
-                'device': self.device.pk,
+                'object_id': self.device.pk,
+                'content_type': '33',
                 'role': self.secretrole1.pk,
                 'name': 'Test Secret 4',
                 'plaintext': 'Secret #4 Plaintext',
             },
             {
-                'device': self.device.pk,
+                'object_id': self.device.pk,
+                'content_type': '33',
                 'role': self.secretrole1.pk,
                 'name': 'Test Secret 5',
                 'plaintext': 'Secret #5 Plaintext',
             },
             {
-                'device': self.device.pk,
+                'object_id': self.device.pk,
+                'content_type': '33',
                 'role': self.secretrole1.pk,
                 'name': 'Test Secret 6',
                 'plaintext': 'Secret #6 Plaintext',
@@ -257,7 +261,8 @@ class SecretTest(APITestCase):
     def test_update_secret(self):
 
         data = {
-            'device': self.device.pk,
+            'object_id': self.device.pk,
+            'content_type': '33',
             'role': self.secretrole2.pk,
             'plaintext': 'NewPlaintext',
         }

netbox/secrets/views.py

@@ -1,10 +1,12 @@
 from __future__ import unicode_literals
 
 import base64
+import os
 
 from django.contrib import messages
 from django.contrib.auth.decorators import permission_required, login_required
 from django.contrib.auth.mixins import PermissionRequiredMixin
+from django.contrib.contenttypes.models import ContentType
 from django.db.models import Count
 from django.shortcuts import get_object_or_404, redirect, render
 from django.urls import reverse
@@ -71,7 +73,7 @@ class SecretRoleBulkDeleteView(PermissionRequiredMixin, BulkDeleteView):
 
 @method_decorator(login_required, name='dispatch')
 class SecretListView(ObjectListView):
-    queryset = Secret.objects.select_related('role', 'device')
+    queryset = Secret.objects.select_related('role').prefetch_related('object')
     filter = filters.SecretFilter
     filter_form = forms.SecretFilterForm
     table = tables.SecretTable
@@ -93,11 +95,16 @@ class SecretView(View):
 @permission_required('secrets.add_secret')
 @userkey_required()
 def secret_add(request, pk):
+    path = request.path
+    path = os.path.normpath(path)
+    object = path.split(os.sep)[2]
+    secret = {
+        'devices': lambda pk: Secret(
+            object_id=pk,
+            content_type=ContentType.objects.get(app_label='dcim', model='device')
+        )
+    }[object](pk)
 
-    # Retrieve device
-    device = get_object_or_404(Device, pk=pk)
-
-    secret = Secret(device=device)
     session_key = get_session_key(request)
 
     if request.method == 'POST':
@@ -131,10 +138,14 @@ def secret_add(request, pk):
     else:
         form = forms.SecretForm(instance=secret)
 
+    return_url = {
+        'devices': Device.objects.get(pk=pk).get_absolute_url()
+    }[object]
+
     return render(request, 'secrets/secret_edit.html', {
         'secret': secret,
         'form': form,
-        'return_url': device.get_absolute_url(),
+        'return_url': return_url,
     })
 
 
@@ -247,7 +258,7 @@ class SecretBulkImportView(BulkImportView):
 
 class SecretBulkEditView(PermissionRequiredMixin, BulkEditView):
     permission_required = 'secrets.change_secret'
-    queryset = Secret.objects.select_related('role', 'device')
+    queryset = Secret.objects.select_related('role').prefetch_related('object')
     filter = filters.SecretFilter
     table = tables.SecretTable
     form = forms.SecretBulkEditForm
@@ -256,7 +267,7 @@ class SecretBulkEditView(PermissionRequiredMixin, BulkEditView):
 
 class SecretBulkDeleteView(PermissionRequiredMixin, BulkDeleteView):
     permission_required = 'secrets.delete_secret'
-    queryset = Secret.objects.select_related('role', 'device')
+    queryset = Secret.objects.select_related('role').prefetch_related('object')
     filter = filters.SecretFilter
     table = tables.SecretTable
     default_return_url = 'secrets:secret_list'

netbox/templates/secrets/secret.html

@@ -9,7 +9,7 @@
             <ol class="breadcrumb">
                 <li><a href="{% url 'secrets:secret_list' %}">Secrets</a></li>
                 <li><a href="{% url 'secrets:secret_list' %}?role={{ secret.role.slug }}">{{ secret.role }}</a></li>
-                <li>{{ secret.device }}{% if secret.name %} ({{ secret.name }}){% endif %}</li>
+                <li>{{ secret.object }}{% if secret.name %} ({{ secret.name }}){% endif %}</li>
             </ol>
         </div>
     </div>
@@ -48,9 +48,9 @@
             </div>
             <table class="table table-hover panel-body">
                 <tr>
-                    <td>Device</td>
+                    <td>{{secret.object_label}}</td>
                     <td>
-                        <a href="{% url 'dcim:device' pk=secret.device.pk %}">{{ secret.device }}</a>
+                        <a href="{{ secret.object.get_absolute_url }}">{{ secret.object }}</a>
                     </td>
                 </tr>
                 <tr>

netbox/templates/secrets/secret_edit.html

@@ -21,9 +21,9 @@
                 <div class="panel-heading"><strong>Secret Attributes</strong></div>
                 <div class="panel-body">
                     <div class="form-group">
-                        <label class="col-md-3 control-label required">Device</label>
+                        <label class="col-md-3 control-label required">{{secret.object_label}}</label>
                         <div class="col-md-9">
-                            <p class="form-control-static">{{ secret.device }}</p>
+                            <p class="form-control-static">{{ secret.object }}</p>
                         </div>
                     </div>
                     {% render_field form.role %}