From 81d955ab7de3cccd000fbcbc5bcd60d78134ced2 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@digitalocean.com>
Date: Wed, 3 Aug 2016 12:00:35 -0400
Subject: [PATCH] Rewrote add_available_ipaddresses() to be much more efficient
 and IPv6-friendly

---
 netbox/ipam/tables.py |  8 +++++-
 netbox/ipam/views.py  | 65 ++++++++++++++++++++++++++-----------------
 2 files changed, 46 insertions(+), 27 deletions(-)

diff --git a/netbox/ipam/tables.py b/netbox/ipam/tables.py
index 800c0fb90..4c5bbee3a 100644
--- a/netbox/ipam/tables.py
+++ b/netbox/ipam/tables.py
@@ -43,7 +43,7 @@ IPADDRESS_LINK = """
 {% if record.pk %}
     <a href="{{ record.get_absolute_url }}">{{ record.address }}</a>
 {% elif perms.ipam.add_ipaddress %}
-    <a href="{% url 'ipam:ipaddress_add' %}?address={{ record.1 }}" class="btn btn-xs btn-success">{{ record.0 }} free IP{{ record.0|pluralize }}</a>
+    <a href="{% url 'ipam:ipaddress_add' %}?address={{ record.1 }}" class="btn btn-xs btn-success">{% if record.0 <= 65536 %}{{ record.0 }}{% else %}Lots of{% endif %} free IP{{ record.0|pluralize }}</a>
 {% else %}
     {{ record.0 }}
 {% endif %}
@@ -158,6 +158,9 @@ class PrefixTable(BaseTable):
     class Meta(BaseTable.Meta):
         model = Prefix
         fields = ('pk', 'prefix', 'status', 'vrf', 'tenant', 'site', 'role', 'description')
+        row_attrs = {
+            'class': lambda record: 'success' if not record.pk else '',
+        }
 
 
 class PrefixBriefTable(BaseTable):
@@ -190,6 +193,9 @@ class IPAddressTable(BaseTable):
     class Meta(BaseTable.Meta):
         model = IPAddress
         fields = ('pk', 'address', 'vrf', 'tenant', 'device', 'interface', 'description')
+        row_attrs = {
+            'class': lambda record: 'success' if not isinstance(record, IPAddress) else '',
+        }
 
 
 class IPAddressBriefTable(BaseTable):
diff --git a/netbox/ipam/views.py b/netbox/ipam/views.py
index 86fc53c74..68cab7429 100644
--- a/netbox/ipam/views.py
+++ b/netbox/ipam/views.py
@@ -1,4 +1,4 @@
-from netaddr import IPNetwork, IPSet
+import netaddr
 from django_tables2 import RequestConfig
 
 from django.contrib.auth.mixins import PermissionRequiredMixin
@@ -21,7 +21,7 @@ def add_available_prefixes(parent, prefix_list):
     """
 
     # Find all unallocated space
-    available_prefixes = IPSet(parent) ^ IPSet([p.prefix for p in prefix_list])
+    available_prefixes = netaddr.IPSet(parent) ^ netaddr.IPSet([p.prefix for p in prefix_list])
     available_prefixes = [Prefix(prefix=p) for p in available_prefixes.iter_cidrs()]
 
     # Concatenate and sort complete list of children
@@ -33,37 +33,50 @@ def add_available_prefixes(parent, prefix_list):
 
 def add_available_ipaddresses(prefix, ipaddress_list):
     """
-    Create fake IPAddress objects for all unallocated space within a prefix.
+    Annotate ranges of available IP addresses within a given prefix.
     """
 
-    # Find all unallocated space
-    available_ips = IPSet(prefix) - IPSet([str(ip.address.ip) for ip in ipaddress_list])
-    available_ips = [IPAddress(address=IPNetwork('{}/{}'.format(ip, prefix.prefixlen))) for ip in available_ips]
+    output = []
+    prev_ip = None
+
+    # Determine first and last usable IP
+    if prefix.version == 6 or (prefix.version == 4 and prefix.prefixlen == 31):
+        first_ip_in_prefix = netaddr.IPAddress(prefix.first)
+    else:
+        first_ip_in_prefix = netaddr.IPAddress(prefix.first + 1)
+    if prefix.version == 4 and prefix.prefixlen == 31:
+        last_ip_in_prefix = netaddr.IPAddress(prefix.last)
+    else:
+        last_ip_in_prefix = netaddr.IPAddress(prefix.last - 1)
 
-    # Concatenate and sort complete list of children
-    ipaddress_list = list(ipaddress_list) + available_ips
-    ipaddress_list.sort(key=lambda ip: ip.address)
     if not ipaddress_list:
-        return []
+        return [(
+            int(last_ip_in_prefix - first_ip_in_prefix + 1),
+            '{}/{}'.format(first_ip_in_prefix, prefix.prefixlen)
+        )]
 
-    # Summarize free IPs in the list
-    computed_list = []
-    count = 0
-    prev_ip = ipaddress_list[0]
+    # Account for any available IPs before the first real IP
+    if ipaddress_list[0].address.ip != first_ip_in_prefix:
+        skipped_count = int(ipaddress_list[0].address.ip - first_ip_in_prefix)
+        first_skipped = '{}/{}'.format(first_ip_in_prefix, prefix.prefixlen)
+        output.append((skipped_count, first_skipped))
+
+    # Iterate through existing IPs and annotate free ranges
     for ip in ipaddress_list:
-        if ip.pk:
-            if count:
-                computed_list.append((count, prev_ip))
-                count = 0
-            computed_list.append(ip)
-            continue
-        if not count:
-            prev_ip = ip
-        count += 1
-    if count:
-        computed_list.append((count, prev_ip))
+        if prev_ip:
+            skipped_count = int(ip.address.ip - prev_ip.address.ip)
+            first_skipped = '{}/{}'.format(prev_ip.address.ip + 1, prefix.prefixlen)
+            output.append((skipped_count, first_skipped))
+        output.append(ip)
+        prev_ip = ip
 
-    return computed_list
+    # Include any remaining available IPs
+    if prev_ip.address.ip != last_ip_in_prefix:
+        skipped_count = int(last_ip_in_prefix - prev_ip.address.ip)
+        first_skipped = '{}/{}'.format(prev_ip.address.ip + 1, prefix.prefixlen)
+        output.append((skipped_count, first_skipped))
+
+    return output
 
 
 #