Merge pull request #35 from DanSheps/develop

Fix problems with StaticRoute FilterSet
DanSheps · Sep 16, 2024 · b85c6ac · b85c6ac
2 parents 7a12079 + 9fbcdfa
commit b85c6ac
Showing 1 changed file with 30 additions and 5 deletions.
diff --git a/netbox_routing/filtersets/static.py b/netbox_routing/filtersets/static.py
@@ -2,12 +2,37 @@
 import netaddr
 from django.db.models import Q
 
+from dcim.models import Device
+from ipam.models import VRF
 from netbox.filtersets import NetBoxModelFilterSet
 from netbox_routing.models import StaticRoute
 
 
 class StaticRouteFilterSet(NetBoxModelFilterSet):
 
+    device = django_filters.ModelMultipleChoiceFilter(
+        field_name='devices__name',
+        queryset=Device.objects.all(),
+        to_field_name='name',
+        label='Device (name)',
+    )
+    device_id = django_filters.ModelMultipleChoiceFilter(
+        field_name='devices',
+        queryset=Device.objects.all(),
+        label='Device (ID)',
+    )
+    vrf = django_filters.ModelMultipleChoiceFilter(
+        field_name='vrf__name',
+        queryset=VRF.objects.all(),
+        to_field_name='name',
+        label='VRF (name)',
+    )
+    vrf_id = django_filters.ModelMultipleChoiceFilter(
+        field_name='vrf',
+        queryset=VRF.objects.all(),
+        label='VRF (ID)',
+    )
+
     prefix = django_filters.CharFilter(
         method='filter_prefix',
         label='Prefix',
@@ -20,7 +45,7 @@ class StaticRouteFilterSet(NetBoxModelFilterSet):
 
     class Meta:
         model = StaticRoute
-        fields = ('vrf', 'prefix', 'devices', 'metric', 'next_hop')
+        fields = ('name', 'devices', 'device', 'device_id', 'vrf', 'vrf_id', 'prefix', 'metric', 'next_hop')
 
     def search(self, queryset, name, value):
         if not value.strip():
@@ -40,15 +65,15 @@ def filter_prefix(self, queryset, name, value):
             return queryset
         try:
             query = str(netaddr.IPNetwork(value).cidr)
-            return queryset.filter(prefix=query)
+            return queryset.filter(**{f'{name}': query})
         except (netaddr.AddrFormatError, ValueError):
             return queryset.none()
 
     def filter_address(self, queryset, name, value):
         if not value.strip():
             return queryset
         try:
-            query = str(netaddr.IPAddress(value))
-            return queryset.filter(prefix=query)
-        except (netaddr.AddrFormatError, ValueError):
+            query = netaddr.IPAddress(value)
+            return queryset.filter(**{f'{name}': query})
+        except (netaddr.AddrFormatError, ValueError) as e:
             return queryset.none()