From feb52d5c9342c867f7ade65a233f301137809480 Mon Sep 17 00:00:00 2001
From: jeremystretch <jstretch@netboxlabs.com>
Date: Sat, 25 Feb 2023 16:12:49 -0500
Subject: [PATCH] Add an available-asns API endpoint

---
 netbox/ipam/api/serializers.py | 16 +++++++++
 netbox/ipam/api/urls.py        |  5 +++
 netbox/ipam/api/views.py       | 65 ++++++++++++++++++++++++++++++++++
 netbox/ipam/models/asns.py     | 10 ++++++
 netbox/utilities/constants.py  |  1 +
 5 files changed, 97 insertions(+)

diff --git a/netbox/ipam/api/serializers.py b/netbox/ipam/api/serializers.py
index c155c7663..82e700677 100644
--- a/netbox/ipam/api/serializers.py
+++ b/netbox/ipam/api/serializers.py
@@ -51,6 +51,22 @@ class ASNSerializer(NetBoxModelSerializer):
         ]
 
 
+class AvailableASNSerializer(serializers.Serializer):
+    """
+    Representation of an ASN which does not exist in the database.
+    """
+    asn = serializers.IntegerField(read_only=True)
+
+    def to_representation(self, asn):
+        range = NestedASNRangeSerializer(self.context['range'], context={
+            'request': self.context['request']
+        }).data
+        return {
+            'range': range,
+            'asn': asn,
+        }
+
+
 #
 # VRFs
 #
diff --git a/netbox/ipam/api/urls.py b/netbox/ipam/api/urls.py
index 5e6510d4b..442fd2240 100644
--- a/netbox/ipam/api/urls.py
+++ b/netbox/ipam/api/urls.py
@@ -29,6 +29,11 @@ router.register('l2vpn-terminations', views.L2VPNTerminationViewSet)
 app_name = 'ipam-api'
 
 urlpatterns = [
+    path(
+        'asn-ranges/<int:pk>/available-asns/',
+        views.AvailableASNsView.as_view(),
+        name='asnrange-available-asns'
+    ),
     path(
         'ip-ranges/<int:pk>/available-ips/',
         views.IPRangeAvailableIPAddressesView.as_view(),
diff --git a/netbox/ipam/api/views.py b/netbox/ipam/api/views.py
index a50e61300..92d326e2c 100644
--- a/netbox/ipam/api/views.py
+++ b/netbox/ipam/api/views.py
@@ -209,6 +209,71 @@ def get_results_limit(request):
     return limit
 
 
+class AvailableASNsView(ObjectValidationMixin, APIView):
+    queryset = ASN.objects.all()
+
+    @swagger_auto_schema(responses={200: serializers.AvailableASNSerializer(many=True)})
+    def get(self, request, pk):
+        asnrange = get_object_or_404(ASNRange.objects.restrict(request.user), pk=pk)
+        limit = get_results_limit(request)
+
+        available_asns = asnrange.get_available_asns()[:limit]
+
+        serializer = serializers.AvailableASNSerializer(available_asns, many=True, context={
+            'request': request,
+            'range': asnrange,
+        })
+
+        return Response(serializer.data)
+
+    @swagger_auto_schema(
+        request_body=serializers.AvailableASNSerializer,
+        responses={201: serializers.ASNSerializer(many=True)}
+    )
+    @advisory_lock(ADVISORY_LOCK_KEYS['available-asns'])
+    def post(self, request, pk):
+        self.queryset = self.queryset.restrict(request.user, 'add')
+        asnrange = get_object_or_404(ASNRange.objects.restrict(request.user), pk=pk)
+
+        # Normalize to a list of objects
+        requested_asns = request.data if isinstance(request.data, list) else [request.data]
+
+        # Determine if the requested number of IPs is available
+        available_asns = asnrange.get_available_asns()
+        if len(available_asns) < len(requested_asns):
+            return Response(
+                {
+                    "detail": f"An insufficient number of ASNs are available within {asnrange} "
+                              f"({len(requested_asns)} requested, {len(available_asns)} available)"
+                },
+                status=status.HTTP_409_CONFLICT
+            )
+
+        # Assign ASNs from the list of available IPs and copy VRF assignment from the parent
+        for i, requested_asn in enumerate(requested_asns):
+            requested_asn['asn'] = available_asns[i]
+            requested_asn['range'] = asnrange.pk
+
+        # Initialize the serializer with a list or a single object depending on what was requested
+        context = {'request': request}
+        if isinstance(request.data, list):
+            serializer = serializers.ASNSerializer(data=requested_asns, many=True, context=context)
+        else:
+            serializer = serializers.ASNSerializer(data=requested_asns[0], context=context)
+
+        # Create the new IP address(es)
+        if serializer.is_valid():
+            try:
+                with transaction.atomic():
+                    created = serializer.save()
+                    self._validate_objects(created)
+            except ObjectDoesNotExist:
+                raise PermissionDenied()
+            return Response(serializer.data, status=status.HTTP_201_CREATED)
+
+        return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
+
+
 class AvailablePrefixesView(ObjectValidationMixin, APIView):
     queryset = Prefix.objects.all()
 
diff --git a/netbox/ipam/models/asns.py b/netbox/ipam/models/asns.py
index a2ea59fa6..3ee33286d 100644
--- a/netbox/ipam/models/asns.py
+++ b/netbox/ipam/models/asns.py
@@ -121,3 +121,13 @@ class ASNRange(OrganizationalModel):
     def clean(self):
         if self.end <= self.start:
             raise ValidationError(f"Starting ASN ({self.start}) must be lower than ending ASN ({self.end}).")
+
+    def get_available_asns(self):
+        """
+        Return all available ASNs within this range.
+        """
+        range = set(self.range)
+        existing_asns = set(ASN.objects.filter(range=self).values_list('asn', flat=True))
+        available_asns = sorted(range - existing_asns)
+
+        return available_asns
diff --git a/netbox/utilities/constants.py b/netbox/utilities/constants.py
index 9303e5f3a..096b60a70 100644
--- a/netbox/utilities/constants.py
+++ b/netbox/utilities/constants.py
@@ -43,6 +43,7 @@ ADVISORY_LOCK_KEYS = {
     'available-prefixes': 100100,
     'available-ips': 100200,
     'available-vlans': 100300,
+    'available-asns': 100400,
 }
 
 #