Netbox Script to rebuild device types and modules

rebuild.py

from typing import List, Union

from dcim.models import (
    ConsolePort,
    ConsoleServerPort,
    Device,
    DeviceBay,
    FrontPort,
    Interface,
    Module,
    ModuleBay,
    PowerOutlet,
    PowerPort,
    RearPort,
)
from django.conf import settings

from dcim.models.devices import Device, DeviceType, Module, ModuleType
from extras.scripts import BooleanVar, ChoiceVar, LogLevelChoices, MultiObjectVar, Script


def rebuild_components(script: Script, device_or_module: Union[Device, Module]):
    script.log_debug(
        f"rebuild_device_components(`{script.__class__.__name__}`, [`{device_or_module}`]({device_or_module.get_absolute_url()}))"
    )

    items = [
        ("consoleporttemplates", "consoleports", ConsolePort),
        ("consoleserverporttemplates", "consoleserverports", ConsoleServerPort),
        ("interfacetemplates", "interfaces", Interface),
        ("powerporttemplates", "powerports", PowerPort),
        ("poweroutlettemplates", "poweroutlets", PowerOutlet),
        ("rearporttemplates", "rearports", RearPort),
        ("frontporttemplates", "frontports", FrontPort),
    ]

    if isinstance(device_or_module, Device):
        items.extend(
            [
                ("modulebaytemplates", "modulebays", ModuleBay),
                ("devicebaytemplates", "devicebays", DeviceBay),
            ]
        )
        device_or_module_type = device_or_module.device_type
        template_kwargs = dict(device=device_or_module)
        device = device_or_module
    else:
        device_or_module_type = device_or_module.module_type
        template_kwargs = dict(device=device_or_module.device, module=device_or_module)
        device = device_or_module.device

    for templates, component_attribute, component_model in items:

        templated_objects = component_model.objects.bulk_create(
            [
                template.instantiate(**template_kwargs)
                for template in getattr(device_or_module_type, templates).all()
                if (
                    not hasattr(template, "resolve_name")
                    and template.name not in getattr(device, component_attribute).values_list("name", flat=True)
                )
                or (
                    hasattr(template, "resolve_name")
                    and template.resolve_name(template_kwargs.get("module"))
                    not in getattr(device, component_attribute).values_list("name", flat=True)
                )
            ]
        )
        if len(templated_objects):
            script.log_success(
                f"**[[{device_or_module}]({device_or_module.get_absolute_url()})]** Successfully created {len(templated_objects)} `{component_model.__name__}` objects"
            )
    if isinstance(device_or_module, Device):
        # Avoid bulk_create to handle MPTT
        inventoryitems = []
        for template in device_or_module.device_type.inventoryitemtemplates.all():
            if template.name not in device_or_module.inventoryitems.values_list(
                "name", flat=True
            ):  # pyright: reportGeneralTypeIssues=false
                inventoryitem = template.instantiate(**template_kwargs).save()
                inventoryitems.append(inventoryitem)
        if len(inventoryitems):
            script.log_success(
                f"**[{device_or_module}]** Successfully created {len(inventoryitems)} `InventoryItem` objects"
            )


class DeviceComponents(Script):
    """
    This script checks for missing components on specific Device instances
    """

    class Meta:  # pylint: disable=too-few-public-methods
        """Meta class for setting Script Attributes"""

        name = "Device Components"
        description = "Check for missing components on specific Device instances."
        commit_default = False

    devices = MultiObjectVar(model=Device, required=True)

    loglevel = ChoiceVar(choices=LogLevelChoices, default=LogLevelChoices.LOG_SUCCESS, required=False)

    if settings.DEBUG:
        debug = BooleanVar(description="If checked, the script will wait for a debugger to be attached to the worker")

    def run(self, data: dict, commit: bool) -> str:  # pylint: disable=unused-argument
        # Fields data
        device_qs: List[Device] = data.get("devices", [])

        loglevel = data.get("loglevel", LogLevelChoices.LOG_SUCCESS)
        debug = data.get("debug", False)

        if debug:
            import debugpy  # pylint: disable=import-outside-toplevel

            debugpy.listen(("0.0.0.0", 5678))
            debugpy.wait_for_client()  # blocks execution until client is attached

        for netbox_device in device_qs:
            rebuild_components(self, netbox_device)

        # cleanup logs
        if loglevel is not None:
            loglevel_values = LogLevelChoices.values()
            self.log = [log for log in self.log if log[0] in loglevel_values[loglevel_values.index(loglevel) :]]

        return ""


class DeviceTypeComponents(Script):
    """
    This script checks for missing components on specific DeviceType instances
    """

    class Meta:  # pylint: disable=too-few-public-methods
        """Meta class for setting Script Attributes"""

        name = "Device Type Components"
        description = "Check for missing components on specific DeviceType instances."
        commit_default = False

    device_types = MultiObjectVar(model=DeviceType, required=True)

    loglevel = ChoiceVar(choices=LogLevelChoices, default=LogLevelChoices.LOG_SUCCESS, required=False)

    if settings.DEBUG:
        debug = BooleanVar(description="If checked, the script will wait for a debugger to be attached to the worker")

    def run(self, data: dict, commit: bool) -> str:  # pylint: disable=unused-argument
        # Fields data
        device_type_qs: List[Device] = data.get("device_types", [])

        loglevel = data.get("loglevel", LogLevelChoices.LOG_SUCCESS)
        debug = data.get("debug", False)

        if debug:
            import debugpy  # pylint: disable=import-outside-toplevel

            debugpy.listen(("0.0.0.0", 5678))
            debugpy.wait_for_client()  # blocks execution until client is attached

        for netbox_device in Device.objects.filter(device_type__in=device_type_qs.values_list("pk", flat=True)):  # type: ignore
            rebuild_components(self, netbox_device)

        # cleanup logs
        if loglevel is not None:
            loglevel_values = LogLevelChoices.values()
            self.log = [log for log in self.log if log[0] in loglevel_values[loglevel_values.index(loglevel) :]]

        return ""


class ModuleComponents(Script):
    """
    This script checks for missing components on specific Module instances.

    Existing components will not be adopted, however.
    """

    class Meta:  # pylint: disable=too-few-public-methods
        """Meta class for setting Script Attributes"""

        name = "Module Components"
        description = "Check for missing components on specific Module instances."
        commit_default = False

    modules = MultiObjectVar(model=Module, required=True)

    loglevel = ChoiceVar(choices=LogLevelChoices, default=LogLevelChoices.LOG_SUCCESS, required=False)

    if settings.DEBUG:
        debug = BooleanVar(description="If checked, the script will wait for a debugger to be attached to the worker")

    def run(self, data: dict, commit: bool) -> str:  # pylint: disable=unused-argument
        # Fields data
        modules_qs: List[Module] = data.get("modules", [])

        loglevel = data.get("loglevel", LogLevelChoices.LOG_SUCCESS)
        debug = data.get("debug", False)

        if debug:
            import debugpy  # pylint: disable=import-outside-toplevel

            debugpy.listen(("0.0.0.0", 5678))
            debugpy.wait_for_client()  # blocks execution until client is attached

        for netbox_module in modules_qs:
            rebuild_components(self, netbox_module)

        # cleanup logs
        if loglevel is not None:
            loglevel_values = LogLevelChoices.values()
            self.log = [log for log in self.log if log[0] in loglevel_values[loglevel_values.index(loglevel) :]]

        return ""


class ModuleTypeComponents(Script):
    """
    This script checks for missing components on specific ModuleType instances

    Existing components will not be adopted, however.
    """

    class Meta:  # pylint: disable=too-few-public-methods
        """Meta class for setting Script Attributes"""

        name = "Module Type Components"
        description = "Check for missing components on specific ModuleType instances."
        commit_default = False

    module_types = MultiObjectVar(model=ModuleType, required=True)

    loglevel = ChoiceVar(choices=LogLevelChoices, default=LogLevelChoices.LOG_SUCCESS, required=False)

    if settings.DEBUG:
        debug = BooleanVar(description="If checked, the script will wait for a debugger to be attached to the worker")

    def run(self, data: dict, commit: bool) -> str:  # pylint: disable=unused-argument
        # Fields data
        module_type_qs: List[Module] = data.get("module_types", [])

        loglevel = data.get("loglevel", LogLevelChoices.LOG_SUCCESS)
        debug = data.get("debug", False)

        if debug:
            import debugpy  # pylint: disable=import-outside-toplevel

            debugpy.listen(("0.0.0.0", 5678))
            debugpy.wait_for_client()  # blocks execution until client is attached

        for netbox_module in Module.objects.filter(module_type__in=module_type_qs.values_list("pk", flat=True)):  # type: ignore
            rebuild_components(self, netbox_module)

        # cleanup logs
        if loglevel is not None:
            loglevel_values = LogLevelChoices.values()
            self.log = [log for log in self.log if log[0] in loglevel_values[loglevel_values.index(loglevel) :]]

        return ""