netbox_client.py

# mypy: ignore-errors
"""
This module contains all methods
to communicate with the
NetBox API endpoint: Data Center Infrastructure Main (dcim)
"""
import pynetbox
import pydantic


# Define device models
class Manufacturer(pydantic.BaseModel):
    """
    The device manufacturer.
    Is mandatory to create a device in netbox.
    Is needed to define the device type.
    """
    name: str
    slug: str


class DeviceType(pydantic.BaseModel):
    """
    The device type for a device.
    Is mandatory to create a device.
    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):
    """
    The role of a device in netbox.
    Is mandatory to create a device.
    """
    name: str
    slug: str


class Site(pydantic.BaseModel):
    """
    The site where to place the device.
    Is mandatory to create a device.
    """
    name: str
    slug: str


def connect(api, token):
    """
    Creates a netbox client to communicate
    with the NetBox endpoints
    The responses of the client are returned mostly as Records or RecordSet.
    For more details see the dcim.py file in the
    pynetbox package.

    To create the client a main URL to NetBox is required. Something like:
    api = 'http://127.0.0.1:8001'

    To create a token, login to the NetBox web frontend or
    post request directly a token from the NetBox API:
    http://127.0.0.1:8001/api/users/tokens/provision/
    """
    return pynetbox.api(api, token=token)


def get_all_devices(nbclient):
    return list(nbclient.dcim.devices.all())


# We need sometimes a specific device
# for creating devices
def get_device_by_name(nbclient, device_name):
    device = nbclient.dcim.devices.get(name=device_name)

    if device is None:
        raise ValueError(f"Device: {device_name} not found")
    else:
        return device


# get all interfaces for a device
def get_interfaces_by_device(nbclient, device_name: str, speed: str):
    device = get_device_by_name(nbclient, device_name)

    return list(nbclient.dcim.interfaces.filter(device_id=device.id,
                                                enabled=False,
                                                mark_connected=False,
                                                speed=speed
                                                ))


# Create a interface
def create_interface(nbclient,
                     iface_name: str,
                     type: str,
                     speed: str,
                     device_name: str) -> dict:
    """
    Creates a 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 definidtion have  a look in
    choises.py in the netbox API implementation in module dcim.
    Returns the new interface object as dict.
    """
    device = get_device_by_name(nbclient, device_name)

    new_iface = nbclient.dcim.interfaces.create(name=iface_name,
                                                type=type,
                                                speed=speed,
                                                enabled=False,
                                                mark_connected=False,
                                                device=device.id)

    return dict(new_iface)


def create_device_type(nbclient, manufacturer: str,
                       model: str,
                       slug: str) -> dict:

    # First get manufacturer id
    manufacturer_id = int(nbclient.dcim.manufacturers.get(name=manufacturer).id)

    device_type = DeviceType(**{'manufacturer': manufacturer_id, 'model': model, 'slug': slug})

    # Convert the format from DeviceType to dict
    device_type = dict(device_type)

    # Create device type
    device_type = nbclient.dcim.device_types.create(device_type)
    return dict(device_type)


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


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


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


def create_device(nbclient,
                  fqdn: str,
                  model: str,
                  device_role: str,
                  site: str) -> dict:
    """ Creates a device and
    returns the new device as a dict
    """

    # Get device type id
    device_type = nbclient.dcim.device_types.get(model=model)

    # Get device role id
    device_role = nbclient.dcim.device_roles.get(name=device_role)

    # Get site id
    device_site = nbclient.dcim.sites.get(name=site)

    # Create new device
    new_device = nbclient.dcim.devices.create(name=fqdn,
                                              device_type=device_type.id,
                                              role=device_role.id,
                                              site=device_site.id)

    return dict(new_device)


def attach_interface_to_lag(nbclient, device_name: str, lag_name: str, iface_name: str) -> dict:
    """
    Assign a given interface to a lag.
    Returns the lag object with the assignend interfaces
    """
    # Get device id
    device = get_device_by_name(nbclient, device_name)

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

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

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

    # Update interface
    updated_iface = nbclient.dcim.interfaces.update(iface)

    return dict(updated_iface)


if __name__ == "__main__":
    print(dict(create_device_manufacturer("Juniper", "juniper")))