Merge pull request #85 from ChrisPortman/support-tagging-objects

Support Tagging Objects

src/netbox_initializers/initializers/__init__.py

@@ -3,7 +3,7 @@
 
 from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ObjectDoesNotExist
-from extras.models import CustomField
+from extras.models import CustomField, Tag
 from ruamel.yaml import YAML
 
 INITIALIZER_ORDER = (
@@ -122,6 +122,27 @@ def set_custom_fields_values(self, entity, custom_field_data):
         if save:
             entity.save()
 
+    def set_tags(self, entity, tags):
+        if not tags:
+            return
+
+        if not hasattr(entity, "tags"):
+            raise Exception(f"⚠️ Tags cannot be applied to {entity}'s model")
+
+        ct = ContentType.objects.get_for_model(entity)
+
+        save = False
+        for tag in Tag.objects.filter(name__in=tags):
+            restricted_cts = tag.object_types.all()
+            if restricted_cts and ct not in restricted_cts:
+                raise Exception(f"⚠️ Tag {tag} cannot be applied to {entity}'s model")
+
+            entity.tags.add(tag)
+            save = True
+
+        if save:
+            entity.save()
+
     def split_params(self, params: dict, unique_params: list = None) -> Tuple[dict, dict]:
         """Split params dict into dict with matching params and a dict with default values"""
 

src/netbox_initializers/initializers/aggregates.py

@@ -20,6 +20,7 @@ def load_data(self):
             return
         for params in aggregates:
             custom_field_data = self.pop_custom_fields(params)
+            tags = params.pop("tags", None)
 
             params["prefix"] = IPNetwork(params["prefix"])
 
@@ -45,6 +46,7 @@ def load_data(self):
                 print("🗞️ Created Aggregate", aggregate.prefix)
 
             self.set_custom_fields_values(aggregate, custom_field_data)
+            self.set_tags(aggregate, tags)
 
 
 register_initializer("aggregates", AggregateInitializer)

src/netbox_initializers/initializers/asns.py

@@ -16,6 +16,7 @@ def load_data(self):
         if asns is None:
             return
         for params in asns:
+            tags = params.pop("tags", None)
             for assoc, details in REQUIRED_ASSOCS.items():
                 model, field = details
                 query = {field: params.pop(assoc)}
@@ -35,5 +36,7 @@ def load_data(self):
             if created:
                 print(f"🔡 Created ASN {asn.asn}")
 
+            self.set_tags(asn, tags)
+
 
 register_initializer("asns", ASNInitializer)

src/netbox_initializers/initializers/cables.py

@@ -202,6 +202,8 @@ def load_data(self):
         if cables is None:
             return
         for params in cables:
+            tags = params.pop("tags", None)
+
             params["termination_a_class"] = get_termination_class_by_name(
                 params.get("termination_a_class")
             )
@@ -248,6 +250,7 @@ def load_data(self):
             CableTermination.objects.create(**params_b_term)
 
             print(f"🧷 Created cable {cable} {cable_name}")
+            self.set_tags(cable, tags)
 
 
 register_initializer("cables", CableInitializer)

src/netbox_initializers/initializers/circuit_types.py

@@ -11,6 +11,7 @@ def load_data(self):
         if circuit_types is None:
             return
         for params in circuit_types:
+            tags = params.pop("tags", None)
             custom_field_data = self.pop_custom_fields(params)
 
             matching_params, defaults = self.split_params(params)
@@ -22,6 +23,7 @@ def load_data(self):
                 print("⚡ Created Circuit Type", circuit_type.name)
 
             self.set_custom_fields_values(circuit_type, custom_field_data)
+            self.set_tags(circuit_type, tags)
 
 
 register_initializer("circuit_types", CircuitTypeInitializer)

src/netbox_initializers/initializers/circuits.py

@@ -17,6 +17,7 @@ def load_data(self):
             return
         for params in circuits:
             custom_field_data = self.pop_custom_fields(params)
+            tags = params.pop("tags", None)
 
             for assoc, details in REQUIRED_ASSOCS.items():
                 model, field = details
@@ -38,6 +39,7 @@ def load_data(self):
                 print("⚡ Created Circuit", circuit.cid)
 
             self.set_custom_fields_values(circuit, custom_field_data)
+            self.set_tags(circuit, tags)
 
 
 register_initializer("circuits", CircuitInitializer)

src/netbox_initializers/initializers/cluster_groups.py

@@ -11,13 +11,15 @@ def load_data(self):
         if cluster_groups is None:
             return
         for params in cluster_groups:
+            tags = params.pop("tags", None)
             matching_params, defaults = self.split_params(params)
             cluster_group, created = ClusterGroup.objects.get_or_create(
                 **matching_params, defaults=defaults
             )
 
             if created:
                 print("🗄️ Created Cluster Group", cluster_group.name)
+            self.set_tags(cluster_group, tags)
 
 
 register_initializer("cluster_groups", ClusterGroupInitializer)

src/netbox_initializers/initializers/cluster_types.py

@@ -11,13 +11,15 @@ def load_data(self):
         if cluster_types is None:
             return
         for params in cluster_types:
+            tags = params.pop("tags", None)
             matching_params, defaults = self.split_params(params)
             cluster_type, created = ClusterType.objects.get_or_create(
                 **matching_params, defaults=defaults
             )
 
             if created:
                 print("🧰 Created Cluster Type", cluster_type.name)
+            self.set_tags(cluster_type, tags)
 
 
 register_initializer("cluster_types", ClusterTypesInitializer)

src/netbox_initializers/initializers/clusters.py

@@ -22,6 +22,7 @@ def load_data(self):
             return
         for params in clusters:
             custom_field_data = self.pop_custom_fields(params)
+            tags = params.pop("tags", None)
 
             for assoc, details in REQUIRED_ASSOCS.items():
                 model, field = details
@@ -43,6 +44,7 @@ def load_data(self):
                 print("🗄️ Created cluster", cluster.name)
 
             self.set_custom_fields_values(cluster, custom_field_data)
+            self.set_tags(cluster, tags)
 
 
 register_initializer("clusters", ClusterInitializer)

src/netbox_initializers/initializers/config_templates.py

@@ -22,6 +22,7 @@ def load_data(self):
         if config_templates is None:
             return
         for template in config_templates:
+            tags = template.pop("tags", None)
             matching_params, defaults = self.split_params(template)
             config_template, created = ConfigTemplate.objects.get_or_create(
                 **matching_params, defaults=defaults
@@ -30,6 +31,7 @@ def load_data(self):
             if created:
                 config_template.save()
                 print("🪝 Created Config Template {0}".format(config_template.name))
+            self.set_tags(config_template, tags)
 
 
 register_initializer("config_templates", ConfigTemplateInitializer)

src/netbox_initializers/initializers/contact_groups.py

@@ -14,6 +14,7 @@ def load_data(self):
             return
         for params in contact_groups:
             custom_field_data = self.pop_custom_fields(params)
+            tags = params.pop("tags", None)
 
             for assoc, details in OPTIONAL_ASSOCS.items():
                 if assoc in params:
@@ -31,6 +32,7 @@ def load_data(self):
                 print("🔳 Created Contact Group", contact_group.name)
 
             self.set_custom_fields_values(contact_group, custom_field_data)
+            self.set_tags(contact_group, tags)
 
 
 register_initializer("contact_groups", ContactGroupInitializer)

src/netbox_initializers/initializers/contact_roles.py

@@ -12,6 +12,7 @@ def load_data(self):
             return
         for params in contact_roles:
             custom_field_data = self.pop_custom_fields(params)
+            tags = params.pop("tags", None)
 
             matching_params, defaults = self.split_params(params)
             contact_role, created = ContactRole.objects.get_or_create(
@@ -22,6 +23,7 @@ def load_data(self):
                 print("🔳 Created Contact Role", contact_role.name)
 
             self.set_custom_fields_values(contact_role, custom_field_data)
+            self.set_tags(contact_role, tags)
 
 
 register_initializer("contact_roles", ContactRoleInitializer)

src/netbox_initializers/initializers/contacts.py

@@ -14,6 +14,7 @@ def load_data(self):
             return
         for params in contacts:
             custom_field_data = self.pop_custom_fields(params)
+            tags = params.pop("tags", None)
 
             for assoc, details in OPTIONAL_ASSOCS.items():
                 if assoc in params:
@@ -29,6 +30,7 @@ def load_data(self):
                 print("👩‍💻 Created Contact", contact.name)
 
             self.set_custom_fields_values(contact, custom_field_data)
+            self.set_tags(contact, tags)
 
 
 register_initializer("contacts", ContactInitializer)

src/netbox_initializers/initializers/device_roles.py

@@ -12,6 +12,7 @@ def load_data(self):
         if device_roles is None:
             return
         for params in device_roles:
+            tags = params.pop("tags", None)
 
             if "color" in params:
                 color = params.pop("color")
@@ -28,5 +29,7 @@ def load_data(self):
             if created:
                 print("🎨 Created device role", device_role.name)
 
+            self.set_tags(device_role, tags)
+
 
 register_initializer("device_roles", DeviceRoleInitializer)

src/netbox_initializers/initializers/device_types.py

@@ -85,6 +85,7 @@ def load_data(self):
             return
         for params in device_types:
             custom_field_data = self.pop_custom_fields(params)
+            tags = params.pop("tags", None)
             components = [(v[0], v[1], params.pop(k, [])) for k, v in SUPPORTED_COMPONENTS.items()]
 
             for assoc, details in REQUIRED_ASSOCS.items():
@@ -109,6 +110,7 @@ def load_data(self):
                 print("🔡 Created device type", device_type.manufacturer, device_type.model)
 
             self.set_custom_fields_values(device_type, custom_field_data)
+            self.set_tags(device_type, tags)
 
             for component in components:
                 c_model, c_match_params, c_params = component

src/netbox_initializers/initializers/devices.py

@@ -30,6 +30,7 @@ def load_data(self):
             return
         for params in devices:
             custom_field_data = self.pop_custom_fields(params)
+            tags = params.pop("tags", None)
 
             # primary ips are handled later in `380_primary_ips.py`
             params.pop("primary_ip4", None)
@@ -57,6 +58,7 @@ def load_data(self):
                 print("🖥️  Created device", device.name)
 
             self.set_custom_fields_values(device, custom_field_data)
+            self.set_tags(device, tags)
 
 
 register_initializer("devices", DeviceInitializer)

src/netbox_initializers/initializers/interfaces.py

@@ -24,6 +24,7 @@ def load_data(self):
             return
         for params in interfaces:
             custom_field_data = self.pop_custom_fields(params)
+            tags = params.pop("tags", None)
 
             related_interfaces = {k: params.pop(k, None) for k in RELATED_ASSOCS}
 
@@ -49,6 +50,7 @@ def load_data(self):
                 print(f"🧷 Created interface {interface} on {interface.device}")
 
             self.set_custom_fields_values(interface, custom_field_data)
+            self.set_tags(interface, tags)
 
             for related_field, related_value in related_interfaces.items():
                 if not related_value:

src/netbox_initializers/initializers/ip_addresses.py

@@ -30,6 +30,7 @@ def load_data(self):
             return
         for params in ip_addresses:
             custom_field_data = self.pop_custom_fields(params)
+            tags = params.pop("tags", None)
 
             vm = params.pop("virtual_machine", None)
             device = params.pop("device", None)
@@ -70,6 +71,7 @@ def load_data(self):
                 print("🧬 Created IP Address", ip_address.address)
 
             self.set_custom_fields_values(ip_address, custom_field_data)
+            self.set_tags(ip_address, tags)
 
 
 register_initializer("ip_addresses", IPAddressInitializer)

src/netbox_initializers/initializers/locations.py

@@ -13,6 +13,7 @@ def load_data(self):
         if locations is None:
             return
         for params in locations:
+            tags = params.pop("tags", None)
 
             for assoc, details in OPTIONAL_ASSOCS.items():
                 if assoc in params:
@@ -26,5 +27,7 @@ def load_data(self):
             if created:
                 print("🎨 Created location", location.name)
 
+            self.set_tags(location, tags)
+
 
 register_initializer("locations", LocationInitializer)

src/netbox_initializers/initializers/manufacturers.py

@@ -11,6 +11,7 @@ def load_data(self):
         if manufacturers is None:
             return
         for params in manufacturers:
+            tags = params.pop("tags", None)
             matching_params, defaults = self.split_params(params)
             manufacturer, created = Manufacturer.objects.get_or_create(
                 **matching_params, defaults=defaults
@@ -19,5 +20,7 @@ def load_data(self):
             if created:
                 print("🏭 Created Manufacturer", manufacturer.name)
 
+            self.set_tags(manufacturer, tags)
+
 
 register_initializer("manufacturers", ManufacturerInitializer)

src/netbox_initializers/initializers/platforms.py

@@ -15,6 +15,7 @@ def load_data(self):
         if platforms is None:
             return
         for params in platforms:
+            tags = params.pop("tags", None)
 
             for assoc, details in OPTIONAL_ASSOCS.items():
                 if assoc in params:
@@ -29,5 +30,7 @@ def load_data(self):
             if created:
                 print("💾 Created platform", platform.name)
 
+            self.set_tags(platform, tags)
+
 
 register_initializer("platforms", PlatformInitializer)

src/netbox_initializers/initializers/power_feeds.py

@@ -16,6 +16,7 @@ def load_data(self):
             return
         for params in power_feeds:
             custom_field_data = self.pop_custom_fields(params)
+            tags = params.pop("tags", None)
 
             for assoc, details in REQUIRED_ASSOCS.items():
                 model, field = details
@@ -39,6 +40,7 @@ def load_data(self):
                 print("⚡ Created Power Feed", power_feed.name)
 
             self.set_custom_fields_values(power_feed, custom_field_data)
+            self.set_tags(power_feed, tags)
 
 
 register_initializer("power_feeds", PowerFeedInitializer)

src/netbox_initializers/initializers/power_panels.py

@@ -16,6 +16,7 @@ def load_data(self):
             return
         for params in power_panels:
             custom_field_data = self.pop_custom_fields(params)
+            tags = params.pop("tags", None)
 
             for assoc, details in REQUIRED_ASSOCS.items():
                 model, field = details
@@ -39,6 +40,7 @@ def load_data(self):
                 print("⚡ Created Power Panel", power_panel.site, power_panel.name)
 
             self.set_custom_fields_values(power_panel, custom_field_data)
+            self.set_tags(power_panel, tags)
 
 
 register_initializer("power_panels", PowerPanelInitializer)