From 76236945c84944dd1c4f4f1eb197d2eb8013c3da Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@netboxlabs.com>
Date: Mon, 19 Jun 2023 14:56:43 -0400
Subject: [PATCH] Move get_next_available_prefix()

---
 netbox/ipam/api/views.py | 15 +++------------
 netbox/ipam/utils.py     | 22 +++++++++++++++++++++-
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/netbox/ipam/api/views.py b/netbox/ipam/api/views.py
index fa973ffcb..4c46ecd44 100644
--- a/netbox/ipam/api/views.py
+++ b/netbox/ipam/api/views.py
@@ -14,6 +14,7 @@ from circuits.models import Provider
 from dcim.models import Site
 from ipam import filtersets
 from ipam.models import *
+from ipam.utils import get_next_available_prefix
 from netbox.api.viewsets import NetBoxModelViewSet
 from netbox.api.viewsets.mixins import ObjectValidationMixin
 from netbox.config import get_config
@@ -349,16 +350,6 @@ class AvailableASNsView(AvailableObjectsView):
         return requested_objects
 
 
-# TODO: Move me
-def get_next_available(ipset, prefix_size):
-    for available_prefix in ipset.iter_cidrs():
-        if prefix_size >= available_prefix.prefixlen:
-            allocated_prefix = f"{available_prefix.network}/{prefix_size}"
-            ipset.remove(allocated_prefix)
-            return allocated_prefix
-    return None
-
-
 class AvailablePrefixesView(AvailableObjectsView):
     queryset = Prefix.objects.all()
     read_serializer_class = serializers.AvailablePrefixSerializer
@@ -374,7 +365,7 @@ class AvailablePrefixesView(AvailableObjectsView):
     def check_sufficient_available(self, requested_objects, available_objects):
         available_prefixes = IPSet(available_objects)
         for requested_object in requested_objects:
-            if not get_next_available(available_prefixes, requested_object['prefix_length']):
+            if not get_next_available_prefix(available_prefixes, requested_object['prefix_length']):
                 return False
         return True
 
@@ -389,7 +380,7 @@ class AvailablePrefixesView(AvailableObjectsView):
         for i, request_data in enumerate(requested_objects):
 
             # Find the first available prefix equal to or larger than the requested size
-            if allocated_prefix := get_next_available(available_prefixes, request_data['prefix_length']):
+            if allocated_prefix := get_next_available_prefix(available_prefixes, request_data['prefix_length']):
                 request_data.update({
                     'prefix': allocated_prefix,
                     'vrf': parent.vrf.pk if parent.vrf else None,
diff --git a/netbox/ipam/utils.py b/netbox/ipam/utils.py
index 93a40e5a0..f54c7d41d 100644
--- a/netbox/ipam/utils.py
+++ b/netbox/ipam/utils.py
@@ -1,7 +1,15 @@
 import netaddr
 
 from .constants import *
-from .models import ASN, Prefix, VLAN
+from .models import Prefix, VLAN
+
+__all__ = (
+    'add_available_ipaddresses',
+    'add_available_vlans',
+    'add_requested_prefixes',
+    'get_next_available_prefix',
+    'rebuild_prefixes',
+)
 
 
 def add_requested_prefixes(parent, prefix_list, show_available=True, show_assigned=True):
@@ -184,3 +192,15 @@ def rebuild_prefixes(vrf):
 
     # Final flush of any remaining Prefixes
     Prefix.objects.bulk_update(update_queue, ['_depth', '_children'])
+
+
+def get_next_available_prefix(ipset, prefix_size):
+    """
+    Given a prefix length, allocate the next available prefix from an IPSet.
+    """
+    for available_prefix in ipset.iter_cidrs():
+        if prefix_size >= available_prefix.prefixlen:
+            allocated_prefix = f"{available_prefix.network}/{prefix_size}"
+            ipset.remove(allocated_prefix)
+            return allocated_prefix
+    return None