Update secrets API, views

netbox/secrets/api/serializers.py

@@ -1,10 +1,12 @@
+from django.contrib.contenttypes.models import ContentType
+from drf_yasg.utils import swagger_serializer_method
 from rest_framework import serializers
 
-from dcim.api.nested_serializers import NestedDeviceSerializer
 from extras.api.customfields import CustomFieldModelSerializer
 from extras.api.serializers import TaggedObjectSerializer
+from secrets.constants import SECRET_ASSIGNMENT_MODELS
 from secrets.models import Secret, SecretRole
-from utilities.api import ValidatedModelSerializer
+from utilities.api import ContentTypeField, ValidatedModelSerializer, get_serializer_for_model
 from .nested_serializers import *
 
 
@@ -23,18 +25,27 @@ class SecretRoleSerializer(ValidatedModelSerializer):
 
 class SecretSerializer(TaggedObjectSerializer, CustomFieldModelSerializer):
     url = serializers.HyperlinkedIdentityField(view_name='secrets-api:secret-detail')
-    device = NestedDeviceSerializer()
+    assigned_object_type = ContentTypeField(
+        queryset=ContentType.objects.filter(SECRET_ASSIGNMENT_MODELS)
+    )
+    assigned_object = serializers.SerializerMethodField(read_only=True)
     role = NestedSecretRoleSerializer()
     plaintext = serializers.CharField()
 
     class Meta:
         model = Secret
         fields = [
-            'id', 'url', 'device', 'role', 'name', 'plaintext', 'hash', 'tags', 'custom_fields', 'created',
-            'last_updated',
+            'id', 'url', 'assigned_object_type', 'assigned_object_id', 'assigned_object', 'role', 'name', 'plaintext',
+            'hash', 'tags', 'custom_fields', 'created', 'last_updated',
         ]
         validators = []
 
+    @swagger_serializer_method(serializer_or_field=serializers.DictField)
+    def get_assigned_object(self, obj):
+        serializer = get_serializer_for_model(obj.assigned_object, prefix='Nested')
+        context = {'request': self.context['request']}
+        return serializer(obj.assigned_object, context=context).data
+
     def validate(self, data):
 
         # Encrypt plaintext data using the master key provided from the view context

netbox/secrets/api/views.py

@@ -46,9 +46,7 @@ class SecretRoleViewSet(ModelViewSet):
 #
 
 class SecretViewSet(ModelViewSet):
-    queryset = Secret.objects.prefetch_related(
-        'device__primary_ip4', 'device__primary_ip6', 'role', 'tags',
-    )
+    queryset = Secret.objects.prefetch_related('role', 'tags')
     serializer_class = serializers.SecretSerializer
     filterset_class = filters.SecretFilterSet
 

netbox/secrets/constants.py

@@ -1,5 +1,13 @@
+from django.db.models import Q
+
+
 #
 # Secrets
 #
 
+SECRET_ASSIGNMENT_MODELS = Q(
+    Q(app_label='dcim', model='device') |
+    Q(app_label='virtualization', model='virtualmachine')
+)
+
 SECRET_PLAINTEXT_MAX_LENGTH = 65535

netbox/secrets/forms.py

@@ -1,6 +1,7 @@
 from Crypto.Cipher import PKCS1_OAEP
 from Crypto.PublicKey import RSA
 from django import forms
+from django.contrib.contenttypes.models import ContentType
 
 from dcim.models import Device
 from extras.forms import (
@@ -142,10 +143,11 @@ class SecretForm(BootstrapMixin, CustomFieldModelForm):
 
 
 class SecretCSVForm(CustomFieldModelCSVForm):
-    device = CSVModelChoiceField(
-        queryset=Device.objects.all(),
-        to_field_name='name',
-        help_text='Assigned device'
+    assigned_object_type = CSVModelChoiceField(
+        queryset=ContentType.objects.all(),
+        limit_choices_to=SECRET_ASSIGNMENT_MODELS,
+        to_field_name='model',
+        help_text='Side A type'
     )
     role = CSVModelChoiceField(
         queryset=SecretRole.objects.all(),

netbox/secrets/tests/test_api.py

@@ -80,9 +80,9 @@ class SecretTest(APIViewTestCases.APIViewTestCase):
         SecretRole.objects.bulk_create(secret_roles)
 
         secrets = (
-            Secret(device=device, role=secret_roles[0], name='Secret 1', plaintext='ABC'),
-            Secret(device=device, role=secret_roles[0], name='Secret 2', plaintext='DEF'),
-            Secret(device=device, role=secret_roles[0], name='Secret 3', plaintext='GHI'),
+            Secret(assigned_object=device, role=secret_roles[0], name='Secret 1', plaintext='ABC'),
+            Secret(assigned_object=device, role=secret_roles[0], name='Secret 2', plaintext='DEF'),
+            Secret(assigned_object=device, role=secret_roles[0], name='Secret 3', plaintext='GHI'),
         )
         for secret in secrets:
             secret.encrypt(self.master_key)
@@ -90,19 +90,22 @@ class SecretTest(APIViewTestCases.APIViewTestCase):
 
         self.create_data = [
             {
-                'device': device.pk,
+                'assigned_object_type': 'dcim.device',
+                'assigned_object_id': device.pk,
                 'role': secret_roles[1].pk,
                 'name': 'Secret 4',
                 'plaintext': 'JKL',
             },
             {
-                'device': device.pk,
+                'assigned_object_type': 'dcim.device',
+                'assigned_object_id': device.pk,
                 'role': secret_roles[1].pk,
                 'name': 'Secret 5',
                 'plaintext': 'MNO',
             },
             {
-                'device': device.pk,
+                'assigned_object_type': 'dcim.device',
+                'assigned_object_id': device.pk,
                 'role': secret_roles[1].pk,
                 'name': 'Secret 6',
                 'plaintext': 'PQR',

netbox/secrets/tests/test_views.py

@@ -69,13 +69,14 @@ class SecretTestCase(
 
         # Create one secret per device to allow bulk-editing of names (which must be unique per device/role)
         Secret.objects.bulk_create((
-            Secret(device=devices[0], role=secretroles[0], name='Secret 1', ciphertext=b'1234567890'),
-            Secret(device=devices[1], role=secretroles[0], name='Secret 2', ciphertext=b'1234567890'),
-            Secret(device=devices[2], role=secretroles[0], name='Secret 3', ciphertext=b'1234567890'),
+            Secret(assigned_object=devices[0], role=secretroles[0], name='Secret 1', ciphertext=b'1234567890'),
+            Secret(assigned_object=devices[1], role=secretroles[0], name='Secret 2', ciphertext=b'1234567890'),
+            Secret(assigned_object=devices[2], role=secretroles[0], name='Secret 3', ciphertext=b'1234567890'),
         ))
 
         cls.form_data = {
-            'device': devices[1].pk,
+            'assigned_object_type': 'dcim.device',
+            'assigned_object_id': devices[1].pk,
             'role': secretroles[1].pk,
             'name': 'Secret X',
         }
@@ -100,11 +101,12 @@ class SecretTestCase(
     def test_import_objects(self):
         self.add_permissions('secrets.add_secret')
 
+        device = Device.objects.get(name='Device 1')
         csv_data = (
-            "device,role,name,plaintext",
-            "Device 1,Secret Role 1,Secret 4,abcdefghij",
-            "Device 1,Secret Role 1,Secret 5,abcdefghij",
-            "Device 1,Secret Role 1,Secret 6,abcdefghij",
+            "assigned_object_type,assigned_object_id,role,name,plaintext",
+            f"device,{device.pk},Secret Role 1,Secret 4,abcdefghij",
+            f"device,{device.pk},Secret Role 1,Secret 5,abcdefghij",
+            f"device,{device.pk},Secret Role 1,Secret 6,abcdefghij",
         )
 
         # Set the session_key cookie on the request

netbox/secrets/views.py

@@ -58,7 +58,7 @@ class SecretRoleBulkDeleteView(BulkDeleteView):
 #
 
 class SecretListView(ObjectListView):
-    queryset = Secret.objects.prefetch_related('role', 'device')
+    queryset = Secret.objects.prefetch_related('role', 'tags')
     filterset = filters.SecretFilterSet
     filterset_form = forms.SecretFilterForm
     table = tables.SecretTable
@@ -198,13 +198,13 @@ class SecretBulkImportView(BulkImportView):
 
 
 class SecretBulkEditView(BulkEditView):
-    queryset = Secret.objects.prefetch_related('role', 'device')
+    queryset = Secret.objects.prefetch_related('role')
     filterset = filters.SecretFilterSet
     table = tables.SecretTable
     form = forms.SecretBulkEditForm
 
 
 class SecretBulkDeleteView(BulkDeleteView):
-    queryset = Secret.objects.prefetch_related('role', 'device')
+    queryset = Secret.objects.prefetch_related('role')
     filterset = filters.SecretFilterSet
     table = tables.SecretTable