scripts/nbhelper/tenant.py

#!/usr/bin/env python3

# SPDX-FileCopyrightText: © 2021 Open Networking Foundation <support@opennetworking.org>
# SPDX-License-Identifier: Apache-2.0

# tenant.py
# The tenant abstract object of Netbox Object - Tenant

from .utils import logger, netboxapi
from .device import Device, VirtualMachine
from .network import Prefix


class Tenant:
    def __init__(self):

        from .utils import netboxapi, netbox_config

        self.nbapi = netboxapi
        self.name = netbox_config["tenant_name"]
        self.name_segments = netbox_config.get("prefix_segments", 1)

        self.tenant = self.nbapi.tenancy.tenants.get(name=self.name)

        # Tenant only keep the resources which owns by it
        self.devices = list()
        self.vms = list()
        self.prefixes = dict()

        # Get the Device and Virtual Machines from Netbox API
        for device_data in self.nbapi.dcim.devices.filter(tenant=self.tenant.slug):
            device = Device(device_data)
            device.tenant = self
            self.devices.append(device)

        for vm_data in self.nbapi.virtualization.virtual_machines.filter(
            tenant=self.tenant.slug
        ):
            vm = VirtualMachine(vm_data)
            vm.tenant = self
            self.vms.append(vm)

        vrf = self.nbapi.ipam.vrfs.get(tenant=self.tenant.slug)
        for prefix_data in self.nbapi.ipam.prefixes.filter(vrf_id=vrf.id):
            prefix = Prefix(prefix_data, self.name_segments)
            if prefix_data.description:
                self.prefixes[prefix_data.prefix] = prefix

    def get_prefixes(self):
        """Get the IP Prefixes owns by current tenant"""

        return self.prefixes

    def get_device_by_name(self, name):
        """
        Find the device or VM which belongs to this Tenant,
        If the name wasn't specified, return the management server
        """

        for machine in self.devices + self.vms:
            if name and machine.name == name:
                return machine
            elif machine.data["device_role"]["name"] == "Router":
                return machine

        ret_msg = (
            "The name '%s' wasn't found in this tenant, "
            + "or can't found any Router in this tenant"
        )

        logger.error(ret_msg, name)
        sys.exit(1)

    def get_devices(self, device_types=["server", "router"]):
        """
        Get all devices (Router + Server) belong to this Tenant
        """

        if not device_types:
            return self.devices + self.vms

        ret = []

        for machine in self.devices:
            if machine.data.device_role.slug in device_types:
                ret.append(machine)

        for vm in self.vms:
            if vm.data.role.slug in device_types:
                ret.append(vm)

        return ret