Просмотр исходного кода

#5306: Standardize device/VM assignment for IP addresses & secrets

Jeremy Stretch 5 лет назад
Родитель
Сommit
f504d1aa99
3 измененных файлов с 41 добавлено и 12 удалено
  1. 5 1
      netbox/ipam/forms.py
  2. 32 7
      netbox/secrets/forms.py
  3. 4 4
      netbox/secrets/tests/test_views.py

+ 5 - 1
netbox/ipam/forms.py

@@ -856,6 +856,7 @@ class IPAddressCSVForm(CustomFieldModelCSVForm):
     )
     status = CSVChoiceField(
         choices=IPAddressStatusChoices,
+        required=False,
         help_text='Operational status'
     )
     role = CSVChoiceField(
@@ -888,7 +889,10 @@ class IPAddressCSVForm(CustomFieldModelCSVForm):
 
     class Meta:
         model = IPAddress
-        fields = IPAddress.csv_headers
+        fields = [
+            'address', 'vrf', 'tenant', 'status', 'role', 'device', 'virtual_machine', 'interface', 'is_primary',
+            'dns_name', 'description',
+        ]
 
     def __init__(self, data=None, *args, **kwargs):
         super().__init__(data, *args, **kwargs)

+ 32 - 7
netbox/secrets/forms.py

@@ -143,31 +143,56 @@ class SecretForm(BootstrapMixin, CustomFieldModelForm):
 
 
 class SecretCSVForm(CustomFieldModelCSVForm):
-    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(),
         to_field_name='name',
         help_text='Assigned role'
     )
+    device = CSVModelChoiceField(
+        queryset=Device.objects.all(),
+        required=False,
+        to_field_name='name',
+        help_text='Assigned device'
+    )
+    virtual_machine = CSVModelChoiceField(
+        queryset=VirtualMachine.objects.all(),
+        required=False,
+        to_field_name='name',
+        help_text='Assigned VM'
+    )
     plaintext = forms.CharField(
         help_text='Plaintext secret data'
     )
 
     class Meta:
         model = Secret
-        fields = Secret.csv_headers
+        fields = ['role', 'name', 'plaintext', 'device', 'virtual_machine']
         help_texts = {
             'name': 'Name or username',
         }
 
+    def clean(self):
+        super().clean()
+
+        device = self.cleaned_data.get('device')
+        virtual_machine = self.cleaned_data.get('virtual_machine')
+
+        # Validate device OR VM is assigned
+        if not device and not virtual_machine:
+            raise forms.ValidationError("Secret must be assigned to a device or a virtual machine")
+        if device and virtual_machine:
+            raise forms.ValidationError("Secret cannot be assigned to both a device and a virtual machine")
+
     def save(self, *args, **kwargs):
+
+        # Set device/VM assignment
+        self.instance.assigned_object = self.cleaned_data['device'] or self.cleaned_data['virtual_machine']
+
         s = super().save(*args, **kwargs)
+
+        # Set plaintext on instance
         s.plaintext = str(self.cleaned_data['plaintext'])
+
         return s
 
 

+ 4 - 4
netbox/secrets/tests/test_views.py

@@ -103,10 +103,10 @@ class SecretTestCase(
 
         device = Device.objects.get(name='Device 1')
         csv_data = (
-            "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",
+            "device,role,name,plaintext",
+            f"{device.name},Secret Role 1,Secret 4,abcdefghij",
+            f"{device.name},Secret Role 1,Secret 5,abcdefghij",
+            f"{device.name},Secret Role 1,Secret 6,abcdefghij",
         )
 
         # Set the session_key cookie on the request