netbox_api.py

from collections import Counter, defaultdict
from contextlib import ExitStack
from pathlib import Path

import pynetbox
import requests
import urllib3


class NetBox:
    def __new__(cls, *args, **kwargs):
        return super().__new__(cls)

    def __init__(self, settings):
        self.counter = Counter(
            added=0,
            updated=0,
            manufacturer=0,
            module_added=0,
            module_port_added=0,
            rack_types_added=0,
            images=0,
        )
        self.url = settings.NETBOX_URL
        self.token = settings.NETBOX_TOKEN
        self.handle = settings.handle
        self.netbox = None
        self.ignore_ssl = settings.IGNORE_SSL_ERRORS
        self.connect_api()
        self.existing_manufacturers = self.get_manufacturers()
        self.template_manager = DeviceTypes(self.netbox, self.handle, self.counter, self.ignore_ssl)

    def connect_api(self):
        try:
            self.netbox = pynetbox.api(self.url, token=self.token)
            if self.ignore_ssl:
                self.handle.verbose_log("IGNORE_SSL_ERRORS is True, catching exception and disabling SSL verification.")
                urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
                self.netbox.http_session.verify = False
        except Exception as e:  # noqa: BLE001
            self.handle.exception("Exception", "NetBox API Error", e)

    def get_api(self):
        return self.netbox

    def get_counter(self):
        return self.counter

    def get_manufacturers(self):
        return {item.slug: item for item in self.netbox.dcim.manufacturers.all()}

    def create_manufacturers(self, vendors):
        to_create = []
        self.existing_manufacturers = self.get_manufacturers()
        for vendor in vendors:
            try:
                manufacturer = self.existing_manufacturers[vendor["slug"]]
                self.handle.verbose_log(f"Manufacturer Exists: {manufacturer.name} - {manufacturer.id}")
            except KeyError:
                to_create.append(vendor)
                self.handle.verbose_log(f"Manufacturer queued for addition: {vendor['name']}")

        if to_create:
            try:
                created_manufacturers = self.netbox.dcim.manufacturers.create(to_create)
                for manufacturer in created_manufacturers:
                    self.handle.verbose_log(f"Manufacturer Created: {manufacturer.name} - {manufacturer.id}")
                    self.counter.update({"manufacturer": 1})
            except pynetbox.RequestError as request_error:
                self.handle.log("Error creating manufacturers")
                self.handle.verbose_log(f"Error during manufacturer creation. - {request_error.error}")

    def create_device_types(self, device_types_to_add, replace_existing_images=True):
        for device_type in device_types_to_add:
            # Remove file base path
            src_file = Path(device_type["src"])
            del device_type["src"]

            # Pre-process front/rear_image flag, remove flag if present
            device_images_from_library = {}

            image_base = src_file.parents[2] / "elevation-images" / src_file.parent.name

            for i in ["front_image", "rear_image"]:
                if i in device_type:
                    if device_type[i]:
                        pattern = f"{device_type['slug']}.{i.split('_')[0]}.*"
                        images = list(image_base.glob(pattern))
                        if images:
                            device_images_from_library[i] = images[0]
                        else:
                            self.handle.log(f"Error locating image file using '{image_base / pattern}'")
                    del device_type[i]

            try:
                dt = self.template_manager.existing_device_types[device_type["slug"]]
                self.handle.verbose_log(f"Device Type Exists: {dt.manufacturer.name} - " + f"{dt.model} - {dt.id}")

                # Fix for existing device types with outdated subdevice_role or u_height
                yaml_role = device_type.get("subdevice_role")
                yaml_u_height = device_type.get("u_height")

                needs_update = False

                if yaml_role:
                    current_role = getattr(dt, "subdevice_role", None)
                    current_role_value = current_role.value if current_role else None

                    if current_role_value != yaml_role:
                        dt.subdevice_role = yaml_role
                        needs_update = True

                if yaml_u_height is not None:
                    current_u_height = getattr(dt, "u_height", None)

                    if current_u_height is not None and float(current_u_height) != float(yaml_u_height):
                        dt.u_height = yaml_u_height
                        needs_update = True

                if needs_update:
                    dt.save()
                    self.handle.verbose_log(f"Updated {dt.model}")

            except KeyError:
                try:
                    dt = self.netbox.dcim.device_types.create(device_type)
                    self.template_manager.existing_device_types[dt.slug] = dt
                    self.counter.update({"added": 1})
                    self.handle.verbose_log(f"Device Type Created: {dt.manufacturer.name} - " + f"{dt.model} - {dt.id}")
                except pynetbox.RequestError as e:
                    self.handle.log(
                        f"Error {e.error} creating device type:"
                        f" {device_type['manufacturer']['name']} {device_type['model']}"
                    )
                    continue

            if "interfaces" in device_type:
                self.template_manager.create_interfaces(device_type["interfaces"], dt.id)
            if "power-ports" in device_type:
                self.template_manager.create_power_ports(device_type["power-ports"], dt.id)
            if "power-port" in device_type:
                self.template_manager.create_power_ports(device_type["power-port"], dt.id)
            if "console-ports" in device_type:
                self.template_manager.create_console_ports(device_type["console-ports"], dt.id)
            if "power-outlets" in device_type:
                self.template_manager.create_power_outlets(device_type["power-outlets"], dt.id)
            if "console-server-ports" in device_type:
                self.template_manager.create_console_server_ports(device_type["console-server-ports"], dt.id)
            if "rear-ports" in device_type:
                self.template_manager.create_rear_ports(device_type["rear-ports"], dt.id)
            if "front-ports" in device_type:
                self.template_manager.create_front_ports(device_type["front-ports"], dt.id)
            if "device-bays" in device_type:
                self.template_manager.create_device_bays(device_type["device-bays"], dt.id)
            if "module-bays" in device_type:
                self.template_manager.create_module_bays(device_type["module-bays"], dt.id)

            # Finally, update images if any
            if device_images_from_library:
                self.template_manager.upload_images(
                    self.url,
                    self.token,
                    device_images_from_library,
                    dt.id,
                    replace_existing_images=replace_existing_images,
                )

    def create_rack_types(self, rack_types):
        for curr_mt in rack_types:
            try:
                rack_type_res = self.template_manager.existing_rack_types[curr_mt["slug"]]
                self.handle.verbose_log(
                    f"Rack Type Exists: {rack_type_res.manufacturer.name} - {rack_type_res.model} - {rack_type_res.id}"
                )
            except KeyError:
                try:
                    rack_type_res = self.netbox.dcim.rack_types.create(curr_mt)
                    self.counter.update({"rack_types_added": 1})
                    self.handle.verbose_log(
                        f"Rack Type Created: {rack_type_res.manufacturer.name} - "
                        f"{rack_type_res.model} - {rack_type_res.id}"
                    )
                except pynetbox.RequestError as exce:
                    self.handle.log(f"Error '{exce.error}' creating rack type: " + f"{curr_mt}")

    def create_module_types(self, module_types):
        for curr_mt in module_types:
            if "profile" in curr_mt and isinstance(curr_mt["profile"], str):
                curr_mt["profile"] = {"name": curr_mt["profile"]}

            try:
                module_type_res = self.template_manager.existing_module_types[curr_mt["manufacturer"]["slug"]][
                    curr_mt["model"]
                ]

                self.handle.verbose_log(
                    f"Module Type Exists: {module_type_res.manufacturer.name} - "
                    f"{module_type_res.model} - {module_type_res.id}"
                )
            except KeyError:
                try:
                    module_type_res = self.netbox.dcim.module_types.create(curr_mt)
                    self.counter.update({"module_added": 1})
                    self.handle.verbose_log(
                        f"Module Type Created: {module_type_res.manufacturer.name} - "
                        f"{module_type_res.model} - {module_type_res.id}"
                    )
                except pynetbox.RequestError as exce:
                    self.handle.log(f"Error '{exce.error}' creating module type: " + f"{curr_mt}")
                    continue

            if "interfaces" in curr_mt:
                self.template_manager.create_module_interfaces(curr_mt["interfaces"], module_type_res.id)
            if "power-ports" in curr_mt:
                self.template_manager.create_module_power_ports(curr_mt["power-ports"], module_type_res.id)
            if "console-ports" in curr_mt:
                self.template_manager.create_module_console_ports(curr_mt["console-ports"], module_type_res.id)
            if "power-outlets" in curr_mt:
                self.template_manager.create_module_power_outlets(curr_mt["power-outlets"], module_type_res.id)
            if "console-server-ports" in curr_mt:
                self.template_manager.create_module_console_server_ports(
                    curr_mt["console-server-ports"], module_type_res.id
                )
            if "rear-ports" in curr_mt:
                self.template_manager.create_module_rear_ports(curr_mt["rear-ports"], module_type_res.id)
            if "front-ports" in curr_mt:
                self.template_manager.create_module_front_ports(curr_mt["front-ports"], module_type_res.id)

    def normalize_module_types(self, module_types):
        module_groups = defaultdict(list)
        for mt in module_types:
            manuf_slug = mt["manufacturer"]["slug"]
            model = mt["model"]
            module_groups[(manuf_slug, model)].append(mt)

        for (manuf_slug, model), mt_list in module_groups.items():
            if len(mt_list) > 1:
                for mt in mt_list:
                    pn = mt.get("part_number")
                    if pn:
                        new_model = f"{model} - {pn}"
                        self.handle.log(f"Deduplicating module model: {model} -> {new_model}")
                        mt["model"] = new_model

        return module_types


class DeviceTypes:
    def __new__(cls, *args, **kwargs):
        return super().__new__(cls)

    def __init__(self, netbox, handle, counter, ignore_ssl):
        self.netbox = netbox
        self.handle = handle
        self.counter = counter
        self.ignore_ssl = ignore_ssl

        self.existing_device_types = self.get_device_types()
        self.existing_module_types = self.get_module_types()
        self.existing_rack_types = self.get_rack_types()

    def get_device_types(self):
        return {item.slug: item for item in self.netbox.dcim.device_types.all()}

    def get_rack_types(self):
        return {item.slug: item for item in self.netbox.dcim.rack_types.all()}

    def get_module_types(self):
        # Modules don't have slugs, so we index by manufacturer slug, then model string
        all_module_types = {}
        for item in self.netbox.dcim.module_types.all():
            if item.manufacturer.slug not in all_module_types:
                all_module_types[item.manufacturer.slug] = {}
            all_module_types[item.manufacturer.slug][item.model] = item
        return all_module_types

    def get_power_ports(self, device_type):
        return {str(item): item for item in self.netbox.dcim.power_port_templates.filter(device_type_id=device_type)}

    def get_rear_ports(self, device_type):
        return {str(item): item for item in self.netbox.dcim.rear_port_templates.filter(device_type_id=device_type)}

    def get_module_power_ports(self, module_type):
        return {str(item): item for item in self.netbox.dcim.power_port_templates.filter(module_type_id=module_type)}

    def get_module_rear_ports(self, module_type):
        return {str(item): item for item in self.netbox.dcim.rear_port_templates.filter(module_type_id=module_type)}

    def get_device_type_ports_to_create(self, dcim_ports, device_type, existing_ports):
        to_create = [port for port in dcim_ports if port["name"] not in existing_ports]
        for port in to_create:
            port["device_type"] = device_type

        return to_create

    def get_module_type_ports_to_create(self, module_ports, module_type, existing_ports):
        to_create = [port for port in module_ports if port["name"] not in existing_ports]
        for port in to_create:
            port["module_type"] = module_type

        return to_create

    def create_interfaces(self, interfaces, device_type):
        existing_interfaces = {
            str(item): item for item in self.netbox.dcim.interface_templates.filter(device_type_id=device_type)
        }
        to_create = self.get_device_type_ports_to_create(interfaces, device_type, existing_interfaces)

        if to_create:
            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_device_ports_created(
                            self.netbox.dcim.interface_templates.create(to_create), "Interface"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Interface")

    def create_power_ports(self, power_ports, device_type):
        existing_power_ports = self.get_power_ports(device_type)
        to_create = self.get_device_type_ports_to_create(power_ports, device_type, existing_power_ports)

        if to_create:
            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_device_ports_created(
                            self.netbox.dcim.power_port_templates.create(to_create), "Power Port"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Power Port")

    def create_console_ports(self, console_ports, device_type):
        existing_console_ports = {
            str(item): item for item in self.netbox.dcim.console_port_templates.filter(device_type_id=device_type)
        }
        to_create = self.get_device_type_ports_to_create(console_ports, device_type, existing_console_ports)

        if to_create:
            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_device_ports_created(
                            self.netbox.dcim.console_port_templates.create(to_create), "Console Port"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Console Port")

    def create_power_outlets(self, power_outlets, device_type):
        existing_power_outlets = {
            str(item): item for item in self.netbox.dcim.power_outlet_templates.filter(device_type_id=device_type)
        }
        to_create = self.get_device_type_ports_to_create(power_outlets, device_type, existing_power_outlets)

        if to_create:
            existing_power_ports = self.get_power_ports(device_type)
            for outlet in to_create:
                try:
                    power_port = existing_power_ports[outlet["power_port"]]
                    outlet["power_port"] = power_port.id
                except KeyError:
                    pass

            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_device_ports_created(
                            self.netbox.dcim.power_outlet_templates.create(to_create), "Power Outlet"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Power Outlet")

    def create_console_server_ports(self, console_server_ports, device_type):
        existing_console_server_ports = {
            str(item): item
            for item in self.netbox.dcim.console_server_port_templates.filter(device_type_id=device_type)
        }
        to_create = self.get_device_type_ports_to_create(
            console_server_ports, device_type, existing_console_server_ports
        )

        if to_create:
            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_device_ports_created(
                            self.netbox.dcim.console_server_port_templates.create(to_create), "Console Server Port"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Console Server Port")

    def create_rear_ports(self, rear_ports, device_type):
        existing_rear_ports = self.get_rear_ports(device_type)
        to_create = self.get_device_type_ports_to_create(rear_ports, device_type, existing_rear_ports)

        if to_create:
            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_device_ports_created(
                            self.netbox.dcim.rear_port_templates.create(to_create), "Rear Port"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Rear Port")

    def create_front_ports(self, front_ports, device_type):
        existing_front_ports = {
            str(item): item for item in self.netbox.dcim.front_port_templates.filter(device_type_id=device_type)
        }
        to_create = self.get_device_type_ports_to_create(front_ports, device_type, existing_front_ports)

        if to_create:
            all_rearports = self.get_rear_ports(device_type)
            for port in to_create:
                try:
                    rear_port = all_rearports[port["rear_port"]]
                    port["rear_port"] = rear_port.id
                except KeyError:
                    self.handle.log(
                        f"Could not find Rear Port for Front Port: {port['name']} - {port['type']} - {device_type}"
                    )

            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_device_ports_created(
                            self.netbox.dcim.front_port_templates.create(to_create), "Front Port"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Front Port")

    def create_device_bays(self, device_bays, device_type):
        existing_device_bays = {
            str(item): item for item in self.netbox.dcim.device_bay_templates.filter(device_type_id=device_type)
        }
        to_create = self.get_device_type_ports_to_create(device_bays, device_type, existing_device_bays)

        if to_create:
            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_device_ports_created(
                            self.netbox.dcim.device_bay_templates.create(to_create), "Device Bay"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Device Bay")

    def create_module_bays(self, module_bays, device_type):
        existing_module_bays = {
            str(item): item for item in self.netbox.dcim.module_bay_templates.filter(device_type_id=device_type)
        }
        to_create = self.get_device_type_ports_to_create(module_bays, device_type, existing_module_bays)

        if to_create:
            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_device_ports_created(
                            self.netbox.dcim.module_bay_templates.create(to_create), "Module Bay"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Module Bay")

    def create_module_interfaces(self, module_interfaces, module_type):
        existing_interfaces = {
            str(item): item for item in self.netbox.dcim.interface_templates.filter(module_type_id=module_type)
        }
        to_create = self.get_module_type_ports_to_create(module_interfaces, module_type, existing_interfaces)

        if to_create:
            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_module_ports_created(
                            self.netbox.dcim.interface_templates.create(to_create), "Module Interface"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Module Interface")

    def create_module_power_ports(self, power_ports, module_type):
        existing_power_ports = self.get_module_power_ports(module_type)
        to_create = self.get_module_type_ports_to_create(power_ports, module_type, existing_power_ports)

        if to_create:
            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_module_ports_created(
                            self.netbox.dcim.power_port_templates.create(to_create), "Module Power Port"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Module Power Port")

    def create_module_console_ports(self, console_ports, module_type):
        existing_console_ports = {
            str(item): item for item in self.netbox.dcim.console_port_templates.filter(module_type_id=module_type)
        }
        to_create = self.get_module_type_ports_to_create(console_ports, module_type, existing_console_ports)

        if to_create:
            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_module_ports_created(
                            self.netbox.dcim.console_port_templates.create(to_create), "Module Console Port"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Module Console Port")

    def create_module_power_outlets(self, power_outlets, module_type):
        existing_power_outlets = {
            str(item): item for item in self.netbox.dcim.power_outlet_templates.filter(module_type_id=module_type)
        }
        to_create = self.get_module_type_ports_to_create(power_outlets, module_type, existing_power_outlets)

        if to_create:
            existing_power_ports = self.get_module_power_ports(module_type)
            for outlet in to_create:
                try:
                    power_port = existing_power_ports[outlet["power_port"]]
                    outlet["power_port"] = power_port.id
                except KeyError:
                    pass

            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_module_ports_created(
                            self.netbox.dcim.power_outlet_templates.create(to_create), "Module Power Outlet"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Module Power Outlet")

    def create_module_console_server_ports(self, console_server_ports, module_type):
        existing_console_server_ports = {
            str(item): item
            for item in self.netbox.dcim.console_server_port_templates.filter(module_type_id=module_type)
        }
        to_create = self.get_module_type_ports_to_create(
            console_server_ports, module_type, existing_console_server_ports
        )

        if to_create:
            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_module_ports_created(
                            self.netbox.dcim.console_server_port_templates.create(to_create),
                            "Module Console Server Port",
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Module Console Server Port")

    def create_module_rear_ports(self, rear_ports, module_type):
        existing_rear_ports = self.get_module_rear_ports(module_type)
        to_create = self.get_module_type_ports_to_create(rear_ports, module_type, existing_rear_ports)

        if to_create:
            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_module_ports_created(
                            self.netbox.dcim.rear_port_templates.create(to_create), "Module Rear Port"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Module Rear Port")

    def create_module_front_ports(self, front_ports, module_type):
        existing_front_ports = {
            str(item): item for item in self.netbox.dcim.front_port_templates.filter(module_type_id=module_type)
        }
        to_create = self.get_module_type_ports_to_create(front_ports, module_type, existing_front_ports)

        if to_create:
            existing_rear_ports = self.get_module_rear_ports(module_type)
            for port in to_create:
                try:
                    rear_port = existing_rear_ports[port["rear_port"]]
                    port["rear_port"] = rear_port.id
                except KeyError:
                    self.handle.log(
                        f"Could not find Rear Port for Front Port: {port['name']} - {port['type']} - {module_type}"
                    )

            try:
                self.counter.update(
                    {
                        "updated": self.handle.log_module_ports_created(
                            self.netbox.dcim.front_port_templates.create(to_create), "Module Front Port"
                        )
                    }
                )
            except pynetbox.RequestError as excep:
                self.handle.log(f"Error '{excep.error}' creating Module Front Port")

    def upload_images(self, baseurl, token, images, device_type_id, replace_existing_images=False):
        """Upload front_image and/or rear_image for the given device type if they do not already exist, unless replace_existing_images is True.

        Args:
        baseurl: URL for Netbox instance
        token: Token to access Netbox instance
        images: map of front_image and/or rear_image filename
        device_type_id: id for the device-type to update
        replace_existing_images: boolean flag, if True, forces upload even if images already exist
        """
        # Fetch the current device type details using PyNetBox
        device_type = self.netbox.dcim.device_types.get(device_type_id)

        if not device_type:
            self.handle.log(f"Device type with ID {device_type_id} not found in NetBox.")
            return

        # Check if images already exist
        existing_images = {
            "front_image": getattr(device_type, "front_image", None),
            "rear_image": getattr(device_type, "rear_image", None),
        }

        # Filter out images that already exist unless replace_existing_images is True
        images_to_upload = {
            key: value for key, value in images.items() if replace_existing_images or not existing_images.get(key)
        }

        if not images_to_upload:
            self.handle.verbose_log(f"No new images to upload for device type {device_type}. Skipping.")
            return

        url = f"{baseurl}/api/dcim/device-types/{device_type_id}/"
        headers = {"Authorization": f"Token {token}"}

        with ExitStack() as stack:
            files = {
                field_name: (path_obj.name, stack.enter_context(path_obj.open("rb")))
                for field_name, path_obj in images_to_upload.items()
            }

            response = requests.patch(url, headers=headers, files=files, verify=(not self.ignore_ssl))

        self.handle.verbose_log(f"Images {images} updated at {url}: {response}")
        self.counter["images"] += len(images)