Refactor nbhelper
Change-Id: I69d10d164fac3eb319e072447a520905880c31dd
diff --git a/scripts/nbhelper/tenant.py b/scripts/nbhelper/tenant.py
new file mode 100644
index 0000000..4d529a7
--- /dev/null
+++ b/scripts/nbhelper/tenant.py
@@ -0,0 +1,92 @@
+#!/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