netbox-plugin-services/netbox_services/engine/allocators.py

"""Allocators create or acquire objects that don't exist yet.

Strategies:
  - ``create``     — instantiate a new model instance from attrs (unsaved)
  - ``first_free`` — allocate the first available IP from a prefix
  - ``specific``   — use an existing object by address or context path
"""

from django.db.models import Model

from .registry import register_allocator, get_selector
from .templating import render_attrs


@register_allocator("create")
class CreateAllocator:
    """Instantiate a new model instance from a dict of attrs.

    The instance is returned **unsaved** — the resolver saves it after
    applying the parent link (if any).

    Attrs values may be:
      - Plain literals (int, bool, list) — passed through
      - Jinja2 strings with ``{{ }}`` — rendered with ``service``,
        ``index``, ``p1`` (parent), ``p2`` (grandparent), etc.
      - FK fields: the rendered value is treated as a **PK** and
        resolved to the related model instance via ``.objects.get(pk=...)``

    Spec shape::

        {"strategy": "create", "attrs": {"name": "rtr-{{ service.tenant.slug }}-{{ index + 1 }}", "site": "{{ ..site.pk }}", "status": "active"}}
    """

    def __init__(self, spec):
        self.spec = spec
        self.attrs = spec.get("attrs", {})

    def allocate(self, context, model=None, index=0, ancestors=None):
        if model is None:
            raise ValueError("CreateAllocator requires a model class.")

        rendered = render_attrs(
            self.attrs, context.service, ancestors or [], index,
        )

        kwargs = {}
        for key, value in rendered.items():
            try:
                field = model._meta.get_field(key)
            except Exception:
                kwargs[key] = value
                continue

            if (
                field.is_relation
                and field.many_to_one
                and not isinstance(value, Model)
                and value is not None
            ):
                related = field.remote_field.model
                kwargs[key] = related.objects.get(pk=int(value))
            else:
                kwargs[key] = value

        return model(**kwargs)


@register_allocator("first_free")
class FirstFreeIPAddressAllocator:
    """Allocate the first available IP address from a resolved prefix.

    Spec shape::

        {"strategy": "first_free", "prefix": {"type": "reference", "path": "parent.site.mgmt_prefix"}}
        {"strategy": "first_free", "prefix": {"type": "prefix", "filters": {"site": "hq"}}}
    """

    def __init__(self, spec):
        self.spec = spec
        self.prefix_spec = spec.get("prefix", {})

    def resolve_prefix(self, context):
        from ipam.models import Prefix
        selector = get_selector(self.prefix_spec)
        qs = selector.resolve(context, Prefix)
        return qs.first()

    def allocate(self, context, model=None, index=0, ancestors=None):
        from ipam.models import IPAddress
        from ipam.choices import IPAddressStatusChoices

        prefix = self.resolve_prefix(context)
        if prefix is None:
            raise ValueError("No prefix resolved for first_free allocation.")

        available = prefix.get_first_available_ip()
        if available is None:
            raise ValueError(
                f"No available IP addresses in prefix {prefix.prefix}."
            )

        ip = IPAddress.objects.create(
            address=available,
            status=IPAddressStatusChoices.STATUS_ACTIVE,
        )
        return ip


@register_allocator("specific")
class SpecificAllocator:
    """Use a specific object identified by content type + ID or context path.

    Spec shape (by reference)::

        {"strategy": "specific", "path": "parent.site.primary_ip4"}

    Spec shape (by value, for IP addresses)::

        {"strategy": "specific", "address": "10.0.0.1/24"}
    """

    def __init__(self, spec):
        self.spec = spec
        self.path = spec.get("path")
        self.address = spec.get("address")

    def allocate(self, context, model=None, index=0, ancestors=None):
        if self.address:
            from ipam.models import IPAddress
            from ipam.choices import IPAddressStatusChoices
            ip, _ = IPAddress.objects.get_or_create(
                address=self.address,
                defaults={"status": IPAddressStatusChoices.STATUS_ACTIVE},
            )
            return ip

        if self.path:
            return context.resolve_path(self.path)

        raise ValueError("SpecificAllocator requires 'path' or 'address'.")