From 9874cef369ba91e26e09b4151b704e20c22480c2 Mon Sep 17 00:00:00 2001
From: Tobias Genannt <tobias.genannt@gmail.com>
Date: Wed, 7 Apr 2021 15:18:33 +0200
Subject: [PATCH] VLAN Groups can be scoped to multiple types

---
 initializers/vlan_groups.yml       | 22 ++++++++++++++++++++--
 startup_scripts/200_vlan_groups.py | 18 +++++++++++++-----
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/initializers/vlan_groups.yml b/initializers/vlan_groups.yml
index 3e0bb00..a40fd0a 100644
--- a/initializers/vlan_groups.yml
+++ b/initializers/vlan_groups.yml
@@ -1,6 +1,24 @@
 # - name: VLAN group 1
-#   site: AMS 1
+#   scope_type: "dcim.region"
+#   scope: Amsterdam
 #   slug: vlan-group-1
 # - name: VLAN group 2
-#   site: AMS 1
+#   scope_type: "dcim.site"
+#   scope: AMS 1
 #   slug: vlan-group-2
+# - name: VLAN group 3
+#   scope_type: "dcim.location"
+#   scope: cage 101
+#   slug: vlan-group-3
+# - name: VLAN group 4
+#   scope_type: "dcim.rack"
+#   scope: rack-01
+#   slug: vlan-group-4
+# - name: VLAN group 5
+#   scope_type: "virtualization.cluster"
+#   scope: cluster1
+#   slug: vlan-group-5
+# - name: VLAN group 6
+#   scope_type: "virtualization.clustergroup"
+#   scope: Group 1
+#   slug: vlan-group-6
diff --git a/startup_scripts/200_vlan_groups.py b/startup_scripts/200_vlan_groups.py
index 7b72b2d..6dda214 100644
--- a/startup_scripts/200_vlan_groups.py
+++ b/startup_scripts/200_vlan_groups.py
@@ -1,6 +1,6 @@
 import sys
 
-from dcim.models import Site
+from django.contrib.contenttypes.models import ContentType
 from ipam.models import VLANGroup
 from startup_script_utils import load_yaml, pop_custom_fields, set_custom_fields_values
 
@@ -9,7 +9,7 @@ vlan_groups = load_yaml("/opt/netbox/initializers/vlan_groups.yml")
 if vlan_groups is None:
     sys.exit()
 
-optional_assocs = {"site": (Site, "name")}
+optional_assocs = {"scope": (None, "name")}
 
 for params in vlan_groups:
     custom_field_data = pop_custom_fields(params)
@@ -18,9 +18,17 @@ for params in vlan_groups:
         if assoc in params:
             model, field = details
             query = {field: params.pop(assoc)}
-
-            params[assoc] = model.objects.get(**query)
-
+            # Get model from Contenttype
+            scope_type = params.pop("scope_type", None)
+            if not scope_type:
+                print("scope_type is missing from VLAN Group")
+                continue
+            app_label, model = str(scope_type).split(".")
+            ct = ContentType.objects.get(app_label=app_label, model=model)
+            if not ct:
+                print(f"ContentType for app_label = '{app_label}' and model = '{model}' not found")
+                continue
+            params["scope_id"] = ct.model_class().objects.get(**query).id
     vlan_group, created = VLANGroup.objects.get_or_create(**params)
 
     if created: