From 9fae11a42f76e937faf252ab93adc4f61f0a6526 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Mon, 2 Nov 2020 15:39:46 -0500
Subject: [PATCH] Add present_in_vrf filters

---
 netbox/ipam/filters.py            | 19 +++++++++++++++++++
 netbox/ipam/forms.py              | 11 ++++++++---
 netbox/ipam/tests/test_filters.py | 18 ++++++++++++++++++
 3 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/netbox/ipam/filters.py b/netbox/ipam/filters.py
index 988ee86fb..be00d5383 100644
--- a/netbox/ipam/filters.py
+++ b/netbox/ipam/filters.py
@@ -230,6 +230,17 @@ class PrefixFilterSet(BaseFilterSet, TenancyFilterSet, CustomFieldFilterSet, Cre
         to_field_name='rd',
         label='VRF (RD)',
     )
+    present_in_vrf_id = django_filters.ModelChoiceFilter(
+        queryset=VRF.objects.all(),
+        method='filter_present_in_vrf',
+        label='VRF'
+    )
+    present_in_vrf = django_filters.ModelChoiceFilter(
+        queryset=VRF.objects.all(),
+        method='filter_present_in_vrf',
+        to_field_name='rd',
+        label='VRF (RD)',
+    )
     region_id = TreeNodeMultipleChoiceFilter(
         queryset=Region.objects.all(),
         field_name='site__region',
@@ -335,6 +346,14 @@ class PrefixFilterSet(BaseFilterSet, TenancyFilterSet, CustomFieldFilterSet, Cre
         except (AddrFormatError, ValueError):
             return queryset.none()
 
+    def filter_present_in_vrf(self, queryset, name, vrf):
+        if vrf is None:
+            return queryset.none
+        return queryset.filter(
+            Q(vrf=vrf) |
+            Q(vrf__export_targets__in=vrf.import_targets.all())
+        )
+
 
 class IPAddressFilterSet(BaseFilterSet, TenancyFilterSet, CustomFieldFilterSet, CreatedUpdatedFilterSet):
     q = django_filters.CharFilter(
diff --git a/netbox/ipam/forms.py b/netbox/ipam/forms.py
index ea7014121..a31ecdd60 100644
--- a/netbox/ipam/forms.py
+++ b/netbox/ipam/forms.py
@@ -519,8 +519,8 @@ class PrefixBulkEditForm(BootstrapMixin, AddRemoveTagsForm, CustomFieldBulkEditF
 class PrefixFilterForm(BootstrapMixin, TenancyFilterForm, CustomFieldFilterForm):
     model = Prefix
     field_order = [
-        'q', 'within_include', 'family', 'mask_length', 'vrf_id', 'status', 'region', 'site', 'role', 'tenant_group',
-        'tenant', 'is_pool', 'expand',
+        'q', 'within_include', 'family', 'mask_length', 'vrf_id', 'present_in_vrf_id', 'status', 'region', 'site',
+        'role', 'tenant_group', 'tenant', 'is_pool', 'expand',
     ]
     mask_length__lte = forms.IntegerField(
         widget=forms.HiddenInput()
@@ -553,9 +553,14 @@ class PrefixFilterForm(BootstrapMixin, TenancyFilterForm, CustomFieldFilterForm)
     vrf_id = DynamicModelMultipleChoiceField(
         queryset=VRF.objects.all(),
         required=False,
-        label='VRF',
+        label='Assigned VRF',
         null_option='Global'
     )
+    present_in_vrf_id = DynamicModelChoiceField(
+        queryset=VRF.objects.all(),
+        required=False,
+        label='Present in VRF'
+    )
     status = forms.MultipleChoiceField(
         choices=PrefixStatusChoices,
         required=False,
diff --git a/netbox/ipam/tests/test_filters.py b/netbox/ipam/tests/test_filters.py
index fe5ac3280..2fead3d93 100644
--- a/netbox/ipam/tests/test_filters.py
+++ b/netbox/ipam/tests/test_filters.py
@@ -354,12 +354,22 @@ class PrefixTestCase(TestCase):
         )
         Site.objects.bulk_create(sites)
 
+        route_targets = (
+            RouteTarget(name='65000:100'),
+            RouteTarget(name='65000:200'),
+            RouteTarget(name='65000:300'),
+        )
+        RouteTarget.objects.bulk_create(route_targets)
+
         vrfs = (
             VRF(name='VRF 1', rd='65000:100'),
             VRF(name='VRF 2', rd='65000:200'),
             VRF(name='VRF 3', rd='65000:300'),
         )
         VRF.objects.bulk_create(vrfs)
+        vrfs[0].import_targets.add(route_targets[0], route_targets[1], route_targets[2])
+        vrfs[1].export_targets.add(route_targets[1])
+        vrfs[2].export_targets.add(route_targets[2])
 
         vlans = (
             VLAN(vid=1, name='VLAN 1'),
@@ -443,6 +453,14 @@ class PrefixTestCase(TestCase):
         params = {'vrf': [vrfs[0].rd, vrfs[1].rd]}
         self.assertEqual(self.filterset(params, self.queryset).qs.count(), 4)
 
+    def test_present_in_vrf(self):
+        vrf1 = VRF.objects.get(name='VRF 1')
+        vrf2 = VRF.objects.get(name='VRF 2')
+        self.assertEqual(self.filterset({'present_in_vrf_id': vrf1.pk}, self.queryset).qs.count(), 6)
+        self.assertEqual(self.filterset({'present_in_vrf_id': vrf2.pk}, self.queryset).qs.count(), 2)
+        self.assertEqual(self.filterset({'present_in_vrf': vrf1.rd}, self.queryset).qs.count(), 6)
+        self.assertEqual(self.filterset({'present_in_vrf': vrf2.rd}, self.queryset).qs.count(), 2)
+
     def test_region(self):
         regions = Region.objects.all()[:2]
         params = {'region_id': [regions[0].pk, regions[1].pk]}