From db60e8868c0b8618649e23c57be0303c8fe4ca59 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@digitalocean.com>
Date: Wed, 1 Mar 2017 14:23:52 -0500
Subject: [PATCH] Closes #901: Support for filtering prefixes and IP addresses
 by mask length

---
 netbox/ipam/fields.py  |  4 +++-
 netbox/ipam/filters.py | 18 ++++++++++++++++++
 netbox/ipam/forms.py   | 10 +++++++++-
 netbox/ipam/lookups.py | 11 ++++++++++-
 4 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/netbox/ipam/fields.py b/netbox/ipam/fields.py
index 00aeb514b..da07f68d9 100644
--- a/netbox/ipam/fields.py
+++ b/netbox/ipam/fields.py
@@ -6,7 +6,7 @@ from django.db import models
 from .formfields import IPFormField
 from .lookups import (
     EndsWith, IEndsWith, IRegex, IStartsWith, NetContained, NetContainedOrEqual, NetContains, NetContainsOrEquals,
-    NetHost, Regex, StartsWith,
+    NetHost, NetMaskLength, Regex, StartsWith,
 )
 
 
@@ -67,6 +67,7 @@ IPNetworkField.register_lookup(NetContainedOrEqual)
 IPNetworkField.register_lookup(NetContains)
 IPNetworkField.register_lookup(NetContainsOrEquals)
 IPNetworkField.register_lookup(NetHost)
+IPNetworkField.register_lookup(NetMaskLength)
 
 
 class IPAddressField(BaseIPField):
@@ -90,3 +91,4 @@ IPAddressField.register_lookup(NetContainedOrEqual)
 IPAddressField.register_lookup(NetContains)
 IPAddressField.register_lookup(NetContainsOrEquals)
 IPAddressField.register_lookup(NetHost)
+IPAddressField.register_lookup(NetMaskLength)
diff --git a/netbox/ipam/filters.py b/netbox/ipam/filters.py
index 3ed401417..42809b954 100644
--- a/netbox/ipam/filters.py
+++ b/netbox/ipam/filters.py
@@ -92,6 +92,10 @@ class PrefixFilter(CustomFieldFilterSet, django_filters.FilterSet):
         method='search_by_parent',
         label='Parent prefix',
     )
+    mask_length = django_filters.NumberFilter(
+        method='filter_mask_length',
+        label='Mask length',
+    )
     vrf_id = NullableModelMultipleChoiceFilter(
         name='vrf_id',
         queryset=VRF.objects.all(),
@@ -171,6 +175,11 @@ class PrefixFilter(CustomFieldFilterSet, django_filters.FilterSet):
         except AddrFormatError:
             return queryset.none()
 
+    def filter_mask_length(self, queryset, name, value):
+        if not value:
+            return queryset
+        return queryset.filter(prefix__net_mask_length=value)
+
 
 class IPAddressFilter(CustomFieldFilterSet, django_filters.FilterSet):
     q = django_filters.CharFilter(
@@ -181,6 +190,10 @@ class IPAddressFilter(CustomFieldFilterSet, django_filters.FilterSet):
         method='search_by_parent',
         label='Parent prefix',
     )
+    mask_length = django_filters.NumberFilter(
+        method='filter_mask_length',
+        label='Mask length',
+    )
     vrf_id = NullableModelMultipleChoiceFilter(
         name='vrf_id',
         queryset=VRF.objects.all(),
@@ -245,6 +258,11 @@ class IPAddressFilter(CustomFieldFilterSet, django_filters.FilterSet):
         except AddrFormatError:
             return queryset.none()
 
+    def filter_mask_length(self, queryset, name, value):
+        if not value:
+            return queryset
+        return queryset.filter(address__net_mask_length=value)
+
 
 class VLANGroupFilter(django_filters.FilterSet):
     site_id = NullableModelMultipleChoiceFilter(
diff --git a/netbox/ipam/forms.py b/netbox/ipam/forms.py
index dd215fb16..681a3c3a0 100644
--- a/netbox/ipam/forms.py
+++ b/netbox/ipam/forms.py
@@ -21,6 +21,12 @@ IP_FAMILY_CHOICES = [
     (6, 'IPv6'),
 ]
 
+PREFIX_MASK_LENGTH_CHOICES = [
+    ('', '---------'),
+] + [(i, i) for i in range(1, 128)]
+
+IPADDRESS_MASK_LENGTH_CHOICES = PREFIX_MASK_LENGTH_CHOICES + [(128, 128)]
+
 
 #
 # VRFs
@@ -266,6 +272,7 @@ class PrefixFilterForm(BootstrapMixin, CustomFieldFilterForm):
         'placeholder': 'Prefix',
     }))
     family = forms.ChoiceField(required=False, choices=IP_FAMILY_CHOICES, label='Address family')
+    mask_length = forms.ChoiceField(required=False, choices=PREFIX_MASK_LENGTH_CHOICES, label='Mask length')
     vrf = FilterChoiceField(
         queryset=VRF.objects.annotate(filter_count=Count('prefixes')),
         to_field_name='rd',
@@ -503,7 +510,8 @@ class IPAddressFilterForm(BootstrapMixin, CustomFieldFilterForm):
     parent = forms.CharField(required=False, label='Parent Prefix', widget=forms.TextInput(attrs={
         'placeholder': 'Prefix',
     }))
-    family = forms.ChoiceField(required=False, choices=IP_FAMILY_CHOICES, label='Address Family')
+    family = forms.ChoiceField(required=False, choices=IP_FAMILY_CHOICES, label='Address family')
+    mask_length = forms.ChoiceField(required=False, choices=IPADDRESS_MASK_LENGTH_CHOICES, label='Mask length')
     vrf = FilterChoiceField(
         queryset=VRF.objects.annotate(filter_count=Count('ip_addresses')),
         to_field_name='rd',
diff --git a/netbox/ipam/lookups.py b/netbox/ipam/lookups.py
index b29890723..8346169ce 100644
--- a/netbox/ipam/lookups.py
+++ b/netbox/ipam/lookups.py
@@ -1,4 +1,4 @@
-from django.db.models import Lookup
+from django.db.models import Lookup, Transform, IntegerField
 from django.db.models.lookups import BuiltinLookup
 
 
@@ -87,3 +87,12 @@ class NetHost(Lookup):
             rhs_params[0] = rhs_params[0].split('/')[0]
         params = lhs_params + rhs_params
         return 'HOST(%s) = %s' % (lhs, rhs), params
+
+
+class NetMaskLength(Transform):
+    lookup_name = 'net_mask_length'
+    function = 'MASKLEN'
+
+    @property
+    def output_field(self):
+        return IntegerField()