From b2a14d4654d10c55d18c080470cadee17f1f9e11 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Fri, 18 Sep 2020 16:23:17 -0400
Subject: [PATCH] Extend secret filters

---
 netbox/secrets/filters.py            | 23 +++++++++++++++++++++
 netbox/secrets/tests/test_filters.py | 30 +++++++++++++++++++++++-----
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/netbox/secrets/filters.py b/netbox/secrets/filters.py
index 9ccc86de4..003859b2f 100644
--- a/netbox/secrets/filters.py
+++ b/netbox/secrets/filters.py
@@ -4,6 +4,7 @@ from django.db.models import Q
 from dcim.models import Device
 from extras.filters import CustomFieldFilterSet, CreatedUpdatedFilterSet
 from utilities.filters import BaseFilterSet, NameSlugSearchFilterSet, TagFilter
+from virtualization.models import VirtualMachine
 from .models import Secret, SecretRole
 
 
@@ -35,6 +36,28 @@ class SecretFilterSet(BaseFilterSet, CustomFieldFilterSet, CreatedUpdatedFilterS
         to_field_name='slug',
         label='Role (slug)',
     )
+    device = django_filters.ModelMultipleChoiceFilter(
+        field_name='device__name',
+        queryset=Device.objects.all(),
+        to_field_name='name',
+        label='Device (name)',
+    )
+    device_id = django_filters.ModelMultipleChoiceFilter(
+        field_name='device',
+        queryset=Device.objects.all(),
+        label='Device (ID)',
+    )
+    virtual_machine = django_filters.ModelMultipleChoiceFilter(
+        field_name='virtual_machine__name',
+        queryset=VirtualMachine.objects.all(),
+        to_field_name='name',
+        label='Virtual machine (name)',
+    )
+    virtual_machine_id = django_filters.ModelMultipleChoiceFilter(
+        field_name='virtual_machine',
+        queryset=VirtualMachine.objects.all(),
+        label='Virtual machine (ID)',
+    )
     tag = TagFilter()
 
     class Meta:
diff --git a/netbox/secrets/tests/test_filters.py b/netbox/secrets/tests/test_filters.py
index b7ac73f1d..0be1ef594 100644
--- a/netbox/secrets/tests/test_filters.py
+++ b/netbox/secrets/tests/test_filters.py
@@ -3,6 +3,7 @@ from django.test import TestCase
 from dcim.models import Device, DeviceRole, DeviceType, Manufacturer, Site
 from secrets.filters import *
 from secrets.models import Secret, SecretRole
+from virtualization.models import Cluster, ClusterType, VirtualMachine
 
 
 class SecretRoleTestCase(TestCase):
@@ -51,6 +52,15 @@ class SecretTestCase(TestCase):
         )
         Device.objects.bulk_create(devices)
 
+        cluster_type = ClusterType.objects.create(name='Cluster Type 1', slug='cluster-type-1')
+        cluster = Cluster.objects.create(name='Cluster 1', type=cluster_type)
+        virtual_machines = (
+            VirtualMachine(name='Virtual Machine 1', cluster=cluster),
+            VirtualMachine(name='Virtual Machine 2', cluster=cluster),
+            VirtualMachine(name='Virtual Machine 3', cluster=cluster),
+        )
+        VirtualMachine.objects.bulk_create(virtual_machines)
+
         roles = (
             SecretRole(name='Secret Role 1', slug='secret-role-1'),
             SecretRole(name='Secret Role 2', slug='secret-role-2'),
@@ -59,9 +69,12 @@ class SecretTestCase(TestCase):
         SecretRole.objects.bulk_create(roles)
 
         secrets = (
-            Secret(device=devices[0], role=roles[0], name='Secret 1', plaintext='SECRET DATA'),
-            Secret(device=devices[1], role=roles[1], name='Secret 2', plaintext='SECRET DATA'),
-            Secret(device=devices[2], role=roles[2], name='Secret 3', plaintext='SECRET DATA'),
+            Secret(assigned_object=devices[0], role=roles[0], name='Secret 1', plaintext='SECRET DATA'),
+            Secret(assigned_object=devices[1], role=roles[1], name='Secret 2', plaintext='SECRET DATA'),
+            Secret(assigned_object=devices[2], role=roles[2], name='Secret 3', plaintext='SECRET DATA'),
+            Secret(assigned_object=virtual_machines[0], role=roles[0], name='Secret 4', plaintext='SECRET DATA'),
+            Secret(assigned_object=virtual_machines[1], role=roles[1], name='Secret 5', plaintext='SECRET DATA'),
+            Secret(assigned_object=virtual_machines[2], role=roles[2], name='Secret 6', plaintext='SECRET DATA'),
         )
         # Must call save() to encrypt Secrets
         for s in secrets:
@@ -78,9 +91,9 @@ class SecretTestCase(TestCase):
     def test_role(self):
         roles = SecretRole.objects.all()[:2]
         params = {'role_id': [roles[0].pk, roles[1].pk]}
-        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 4)
         params = {'role': [roles[0].slug, roles[1].slug]}
-        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 4)
 
     def test_device(self):
         devices = Device.objects.all()[:2]
@@ -88,3 +101,10 @@ class SecretTestCase(TestCase):
         self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
         params = {'device': [devices[0].name, devices[1].name]}
         self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+
+    def test_virtual_machine(self):
+        virtual_machines = VirtualMachine.objects.all()[:2]
+        params = {'virtual_machine_id': [virtual_machines[0].pk, virtual_machines[1].pk]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)
+        params = {'virtual_machine': [virtual_machines[0].name, virtual_machines[1].name]}
+        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)