netbox_client.py

"""Contain all methods to communicate with the NetBox API endpoint. Data Center Infrastructure Main (DCIM)."""
from uuid import UUID

import pydantic
import pynetbox
from pynetbox.models.dcim import Devices, DeviceTypes, Interfaces

from gso.products import Router
from gso.settings import load_oss_params
from gso.utils.device_info import FEASIBLE_IP_TRUNK_LAG_RANGE, TierInfo
from gso.utils.exceptions import NotFoundError, WorkflowStateError


# Define device models
class Manufacturer(pydantic.BaseModel):
    """Defines the manufacturer of a device."""

    name: str
    slug: str


class DeviceType(pydantic.BaseModel):
    """Defines the device type.

    The manufacturer should be created first to get the manufacturer id,
    which is defined here as int.
    """

    manufacturer: int
    model: str
    slug: str


class DeviceRole(pydantic.BaseModel):
    """Defines the role of a device."""

    name: str
    slug: str


class Site(pydantic.BaseModel):
    """Defines the site of a device."""

    name: str
    slug: str


class NetBoxClient:
    """Implement all methods to communicate with the NetBox API."""

    def __init__(self) -> None:
        self.netbox_params = load_oss_params().NETBOX
        self.netbox = pynetbox.api(self.netbox_params.api, self.netbox_params.token)

    def get_all_devices(self) -> list[Devices]:
        return list(self.netbox.dcim.devices.all())

    def get_device_by_name(self, device_name: str) -> Devices:
        """Return the device object by name from netbox, or None if not found."""
        return self.netbox.dcim.devices.get(name=device_name)

    def get_interfaces_by_device(self, device_name: str, speed: str) -> list[Interfaces]:
        """Get all interfaces of a device by name and speed."""

        device = self.get_device_by_name(device_name)
        return list(
            self.netbox.dcim.interfaces.filter(device_id=device.id, enabled=False, mark_connected=False, speed=speed)
        )

    def create_interface(self, iface_name: str, type: str, speed: str, device_name: str) -> Interfaces:
        """Create new interface on a device, where device is defined by name.

        The type parameter can be 1000base-t, 10gbase-t, lag...
        For more details on type definition have  a look in
        choices.py in the netbox API implementation in module DCIM.
        Returns the new interface object as dict.
        """
        device = self.get_device_by_name(device_name)
        return self.netbox.dcim.interfaces.create(
            name=iface_name, type=type, speed=speed, enabled=False, mark_connected=False, device=device.id
        )

    def create_device_type(self, manufacturer: str, model: str, slug: str) -> DeviceTypes:
        """Create a new device type in netbox."""

        # First get manufacturer id
        manufacturer_id = int(self.netbox.dcim.manufacturers.get(name=manufacturer).id)
        device_type = DeviceType(**{"manufacturer": manufacturer_id, "model": model, "slug": slug})  # type: ignore
        return self.netbox.dcim.device_types.create(dict(device_type))

    def create_device_role(self, name: str, slug: str) -> DeviceRole:
        device_role = DeviceRole(**{"name": name, "slug": slug})
        return self.netbox.dcim.device_roles.create(dict(device_role))

    def create_device_site(self, name: str, slug: str) -> Site:
        device_site = Site(**{"name": name, "slug": slug})
        return self.netbox.dcim.sites.create(dict(device_site))

    def create_device_manufacturer(self, name: str, slug: str) -> Manufacturer:
        device_manufacturer = Manufacturer(**{"name": name, "slug": slug})
        return self.netbox.dcim.manufacturers.create(dict(device_manufacturer))

    def create_device(self, router_name: str, site_tier: str) -> Devices:
        """Create a new device in netbox."""

        # Get device type id
        tier_info = TierInfo().get_module_by_name(f"Tier{site_tier}")
        device_type = self.netbox.dcim.device_types.get(model=tier_info.device_type)

        # Get device role id
        device_role = self.netbox.dcim.device_roles.get(name="router")

        # Get site id
        device_site = self.netbox.dcim.sites.get(name="Amsterdam")

        # Create new device
        device = self.netbox.dcim.devices.create(
            name=router_name, device_type=device_type.id, role=device_role.id, site=device_site.id
        )
        module_bays = list(self.netbox.dcim.module_bays.filter(device_id=device.id))
        card_type = self.netbox.dcim.module_types.get(model=tier_info.module_type)
        for module_bay in module_bays:
            self.netbox.dcim.modules.create(
                device=device.id,
                module_bay=module_bay.id,
                module_type=card_type.id,
                status="active",
                comments="Installed via pynetbox",
            )

        return device

    def delete_device(self, router_name: str) -> None:
        self.netbox.dcim.devices.get(name=router_name).delete()
        return

    def attach_interface_to_lag(self, device_name: str, lag_name: str, iface_name: str) -> Interfaces:
        """Assign a given interface to a lag.

        Returns the lag object with the assignend interfaces
        """
        # Get device id
        device = self.get_device_by_name(device_name)

        # Now get interface for device
        iface = self.netbox.dcim.interfaces.get(name=iface_name, device_id=device.id)

        # Get lag
        lag = self.netbox.dcim.interfaces.get(name=lag_name, device_id=device.id)

        # Assign interface to lag
        iface.lag = lag.id

        # Update interface
        return self.netbox.dcim.interfaces.update(iface)

    def reserve_interface(self, device_name: str, iface_name: str) -> Interfaces:
        """Reserve an interface by enabling it."""

        # First get interface from device
        device = self.get_device_by_name(device_name)
        interface = self.netbox.dcim.interfaces.get(device_id=device.id, name=iface_name)

        # Reserve interface by enabling it
        if interface is None:
            raise NotFoundError(f"Interface: {iface_name} on device: {device_name} not found.")

        # Check if interface is reserved
        if interface.enabled:
            raise WorkflowStateError(f"The interface: {iface_name} on device: {device_name} is already reserved.")

        # Reserve interface by enabling it
        interface.enabled = True
        interface.save()

        return interface

    def allocate_interface(self, device_name: str, iface_name: str) -> Interfaces:
        """Allocate an interface by marking it as connected."""

        device = self.get_device_by_name(device_name)
        interface = self.netbox.dcim.interfaces.get(device_id=device.id, name=iface_name)

        # Check if interface is available
        if interface is None:
            raise NotFoundError(f"Interface: {iface_name} on device: {device_name} not found.")

        # Check if interface is reserved
        if interface.mark_connected:
            raise WorkflowStateError(f"The interface: {iface_name} on device: {device_name} is already allocated.")

        # allocate interface by mark as connected
        interface.mark_connected = True
        interface.save()

        return interface

    def get_available_lags(self, router_id: UUID) -> list[str]:
        """Return all available lags not assigned to a device."""

        router_name = Router.from_subscription(router_id).router.router_fqdn
        device = self.get_device_by_name(router_name)

        # Get the existing lag interfaces for the device
        lag_interface_names = [
            interface["name"] for interface in self.netbox.dcim.interfaces.filter(device_name=device.id, type="lag")
        ]

        # Generate all feasible lags
        all_feasible_lags = [f"LAG-{i}" for i in FEASIBLE_IP_TRUNK_LAG_RANGE]

        # Return available lags not assigned to the device
        return [lag for lag in all_feasible_lags if lag not in lag_interface_names]