from django.conf import settings from django.core.exceptions import ObjectDoesNotExist, PermissionDenied from django.db import transaction from django.shortcuts import get_object_or_404 from django_pglocks import advisory_lock from drf_yasg.utils import swagger_auto_schema from rest_framework import status from rest_framework.decorators import action from rest_framework.response import Response from ipam.models import * from utilities.constants import ADVISORY_LOCK_KEYS from . import serializers class AvailablePrefixesMixin: @swagger_auto_schema(method='get', responses={200: serializers.AvailablePrefixSerializer(many=True)}) @swagger_auto_schema(method='post', responses={201: serializers.PrefixSerializer(many=False)}) @action(detail=True, url_path='available-prefixes', methods=['get', 'post']) @advisory_lock(ADVISORY_LOCK_KEYS['available-prefixes']) def available_prefixes(self, request, pk=None): """ A convenience method for returning available child prefixes within a parent. The advisory lock decorator uses a PostgreSQL advisory lock to prevent this API from being invoked in parallel, which results in a race condition where multiple insertions can occur. """ prefix = get_object_or_404(self.queryset, pk=pk) available_prefixes = prefix.get_available_prefixes() if request.method == 'POST': # Validate Requested Prefixes' length serializer = serializers.PrefixLengthSerializer( data=request.data if isinstance(request.data, list) else [request.data], many=True, context={ 'request': request, 'prefix': prefix, } ) if not serializer.is_valid(): return Response( serializer.errors, status=status.HTTP_400_BAD_REQUEST ) requested_prefixes = serializer.validated_data # Allocate prefixes to the requested objects based on availability within the parent for i, requested_prefix in enumerate(requested_prefixes): # Find the first available prefix equal to or larger than the requested size for available_prefix in available_prefixes.iter_cidrs(): if requested_prefix['prefix_length'] >= available_prefix.prefixlen: allocated_prefix = '{}/{}'.format(available_prefix.network, requested_prefix['prefix_length']) requested_prefix['prefix'] = allocated_prefix requested_prefix['vrf'] = prefix.vrf.pk if prefix.vrf else None break else: return Response( { "detail": "Insufficient space is available to accommodate the requested prefix size(s)" }, status=status.HTTP_204_NO_CONTENT ) # Remove the allocated prefix from the list of available prefixes available_prefixes.remove(allocated_prefix) # 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.PrefixSerializer(data=requested_prefixes, many=True, context=context) else: serializer = serializers.PrefixSerializer(data=requested_prefixes[0], context=context) # Create the new Prefix(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) else: serializer = serializers.AvailablePrefixSerializer(available_prefixes.iter_cidrs(), many=True, context={ 'request': request, 'vrf': prefix.vrf, }) return Response(serializer.data) class AvailableIPsMixin: parent_model = Prefix @swagger_auto_schema(method='get', responses={200: serializers.AvailableIPSerializer(many=True)}) @swagger_auto_schema(method='post', responses={201: serializers.AvailableIPSerializer(many=True)}, request_body=serializers.AvailableIPSerializer(many=True)) @action(detail=True, url_path='available-ips', methods=['get', 'post'], queryset=IPAddress.objects.all()) @advisory_lock(ADVISORY_LOCK_KEYS['available-ips']) def available_ips(self, request, pk=None): """ A convenience method for returning available IP addresses within a Prefix or IPRange. By default, the number of IPs returned will be equivalent to PAGINATE_COUNT. An arbitrary limit (up to MAX_PAGE_SIZE, if set) may be passed, however results will not be paginated. The advisory lock decorator uses a PostgreSQL advisory lock to prevent this API from being invoked in parallel, which results in a race condition where multiple insertions can occur. """ parent = get_object_or_404(self.parent_model.objects.restrict(request.user), pk=pk) # Create the next available IP if request.method == 'POST': # Normalize to a list of objects requested_ips = request.data if isinstance(request.data, list) else [request.data] # Determine if the requested number of IPs is available available_ips = parent.get_available_ips() if available_ips.size < len(requested_ips): return Response( { "detail": f"An insufficient number of IP addresses are available within {parent} " f"({len(requested_ips)} requested, {len(available_ips)} available)" }, status=status.HTTP_204_NO_CONTENT ) # Assign addresses from the list of available IPs and copy VRF assignment from the parent available_ips = iter(available_ips) for requested_ip in requested_ips: requested_ip['address'] = f'{next(available_ips)}/{parent.mask_length}' requested_ip['vrf'] = parent.vrf.pk if parent.vrf else None # 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.IPAddressSerializer(data=requested_ips, many=True, context=context) else: serializer = serializers.IPAddressSerializer(data=requested_ips[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) # Determine the maximum number of IPs to return else: try: limit = int(request.query_params.get('limit', settings.PAGINATE_COUNT)) except ValueError: limit = settings.PAGINATE_COUNT if settings.MAX_PAGE_SIZE: limit = min(limit, settings.MAX_PAGE_SIZE) # Calculate available IPs within the parent ip_list = [] for index, ip in enumerate(parent.get_available_ips(), start=1): ip_list.append(ip) if index == limit: break serializer = serializers.AvailableIPSerializer(ip_list, many=True, context={ 'request': request, 'parent': parent, }) return Response(serializer.data)