Add function to load YAML files

This commit starts to remove some code redundancy from the startup
scripts for easier maintenance.

Author: tobiasge

startup_scripts/000_users.py

@@ -1,42 +1,37 @@
 from django.contrib.auth.models import Permission, Group, User
 from users.models import Token
 
-from ruamel.yaml import YAML
-from pathlib import Path
+from startup_script_utils import load_yaml
 import sys
 
-file = Path('/opt/netbox/initializers/users.yml')
-if not file.is_file():
+users = load_yaml('/opt/netbox/initializers/users.yml')
+
+if users is None:
   sys.exit()
 
-with file.open('r') as stream:
-  yaml=YAML(typ='safe')
-  users = yaml.load(stream)
-
-  if users is not None:
-    for username, user_details in users.items():
-      if not User.objects.filter(username=username):
-        user = User.objects.create_user(
-          username = username,
-          password = user_details.get('password', 0) or User.objects.make_random_password)
-
-        print("👤 Created user",username)
-
-        if user_details.get('api_token', 0):
-          Token.objects.create(user=user, key=user_details['api_token'])
-
-        yaml_permissions = user_details.get('permissions', [])
-        if yaml_permissions:
-          subject = user.user_permissions
-          subject.clear()
-          for yaml_permission in yaml_permissions:
-            if '*' in yaml_permission:
-              permission_filter = '^' + yaml_permission.replace('*','.*') + '$'
-              permissions = Permission.objects.filter(codename__iregex=permission_filter)
-              print("  ⚿ Granting", permissions.count(), "permissions matching '" + yaml_permission + "'")
-            else:
-              permissions = Permission.objects.filter(codename=yaml_permission)
-              print("  ⚿ Granting permission", yaml_permission)
-            
-            for permission in permissions:
-              subject.add(permission)
+for username, user_details in users.items():
+  if not User.objects.filter(username=username):
+    user = User.objects.create_user(
+      username = username,
+      password = user_details.get('password', 0) or User.objects.make_random_password)
+
+    print("👤 Created user",username)
+
+    if user_details.get('api_token', 0):
+      Token.objects.create(user=user, key=user_details['api_token'])
+
+    yaml_permissions = user_details.get('permissions', [])
+    if yaml_permissions:
+      subject = user.user_permissions
+      subject.clear()
+      for yaml_permission in yaml_permissions:
+        if '*' in yaml_permission:
+          permission_filter = '^' + yaml_permission.replace('*','.*') + '$'
+          permissions = Permission.objects.filter(codename__iregex=permission_filter)
+          print("  ⚿ Granting", permissions.count(), "permissions matching '" + yaml_permission + "'")
+        else:
+          permissions = Permission.objects.filter(codename=yaml_permission)
+          print("  ⚿ Granting permission", yaml_permission)
+
+        for permission in permissions:
+          subject.add(permission)

startup_scripts/010_groups.py

@@ -1,41 +1,35 @@
 from django.contrib.auth.models import Permission, Group, User
-from ruamel.yaml import YAML
-from pathlib import Path
+from startup_script_utils import load_yaml
 import sys
 
-file = Path('/opt/netbox/initializers/groups.yml')
-if not file.is_file():
+groups = load_yaml('/opt/netbox/initializers/groups.yml')
+if groups is None:
   sys.exit()
 
-with file.open('r') as stream:
-  yaml=YAML(typ='safe')
-  groups = yaml.load(stream)
-
-  if groups is not None:
-    for groupname, group_details in groups.items():
-      group, created = Group.objects.get_or_create(name=groupname)
-
-      if created:
-        print("👥 Created group", groupname)
-
-      for username in group_details.get('users', []):
-        user = User.objects.get(username=username)
-
-        if user:
-          user.groups.add(group)
-
-      yaml_permissions = group_details.get('permissions', [])
-      if yaml_permissions:
-        subject = group.permissions
-        subject.clear()
-        for yaml_permission in yaml_permissions:
-          if '*' in yaml_permission:
-            permission_filter = '^' + yaml_permission.replace('*','.*') + '$'
-            permissions = Permission.objects.filter(codename__iregex=permission_filter)
-            print("  ⚿ Granting", permissions.count(), "permissions matching '" + yaml_permission + "'")
-          else:
-            permissions = Permission.objects.filter(codename=yaml_permission)
-            print("  ⚿ Granting permission", yaml_permission)
-
-          for permission in permissions:
-            subject.add(permission)
+for groupname, group_details in groups.items():
+  group, created = Group.objects.get_or_create(name=groupname)
+
+  if created:
+    print("👥 Created group", groupname)
+
+  for username in group_details.get('users', []):
+    user = User.objects.get(username=username)
+
+    if user:
+      user.groups.add(group)
+
+  yaml_permissions = group_details.get('permissions', [])
+  if yaml_permissions:
+    subject = group.permissions
+    subject.clear()
+    for yaml_permission in yaml_permissions:
+      if '*' in yaml_permission:
+        permission_filter = '^' + yaml_permission.replace('*','.*') + '$'
+        permissions = Permission.objects.filter(codename__iregex=permission_filter)
+        print("  ⚿ Granting", permissions.count(), "permissions matching '" + yaml_permission + "'")
+      else:
+        permissions = Permission.objects.filter(codename=yaml_permission)
+        print("  ⚿ Granting permission", yaml_permission)
+
+      for permission in permissions:
+        subject.add(permission)

startup_scripts/020_custom_fields.py

@@ -1,7 +1,6 @@
 from extras.models import CustomField, CustomFieldChoice
 
-from ruamel.yaml import YAML
-from pathlib import Path
+from startup_script_utils import load_yaml
 import sys
 
 def get_class_for_class_path(class_path):
@@ -13,47 +12,43 @@ def get_class_for_class_path(class_path):
   clazz = getattr(module, class_name)
   return ContentType.objects.get_for_model(clazz)
 
-file = Path('/opt/netbox/initializers/custom_fields.yml')
-if not file.is_file():
-  sys.exit()
+customfields = load_yaml('/opt/netbox/initializers/custom_fields.yml')
 
-with file.open('r') as stream:
-  yaml = YAML(typ='safe')
-  customfields = yaml.load(stream)
+if customfields is None:
+  sys.exit()
 
-  if customfields is not None:
-    for cf_name, cf_details in customfields.items():
-      custom_field, created = CustomField.objects.get_or_create(name = cf_name)
+for cf_name, cf_details in customfields.items():
+  custom_field, created = CustomField.objects.get_or_create(name = cf_name)
 
-      if created:
-        if cf_details.get('default', 0):
-          custom_field.default = cf_details['default']
+  if created:
+    if cf_details.get('default', 0):
+      custom_field.default = cf_details['default']
 
-        if cf_details.get('description', 0):
-          custom_field.description = cf_details['description']
+    if cf_details.get('description', 0):
+      custom_field.description = cf_details['description']
 
-        if cf_details.get('label', 0):
-          custom_field.label = cf_details['label']
+    if cf_details.get('label', 0):
+      custom_field.label = cf_details['label']
 
-        for object_type in cf_details.get('on_objects', []):
-          custom_field.obj_type.add(get_class_for_class_path(object_type))
+    for object_type in cf_details.get('on_objects', []):
+      custom_field.obj_type.add(get_class_for_class_path(object_type))
 
-        if cf_details.get('required', 0):
-          custom_field.required = cf_details['required']
+    if cf_details.get('required', 0):
+      custom_field.required = cf_details['required']
 
-        if cf_details.get('type', 0):
-          custom_field.type = cf_details['type']
+    if cf_details.get('type', 0):
+      custom_field.type = cf_details['type']
 
-        if cf_details.get('weight', 0):
-          custom_field.weight = cf_details['weight']
+    if cf_details.get('weight', 0):
+      custom_field.weight = cf_details['weight']
 
-        custom_field.save()
+    custom_field.save()
 
-        for idx, choice_details in enumerate(cf_details.get('choices', [])):
-          choice, _ = CustomFieldChoice.objects.get_or_create(
-            field=custom_field,
-            value=choice_details['value'],
-            defaults={'weight': idx * 10}
-          )
+    for idx, choice_details in enumerate(cf_details.get('choices', [])):
+      choice, _ = CustomFieldChoice.objects.get_or_create(
+        field=custom_field,
+        value=choice_details['value'],
+        defaults={'weight': idx * 10}
+      )
 
-        print("🔧 Created custom field", cf_name)
+    print("🔧 Created custom field", cf_name)

startup_scripts/030_regions.py

@@ -1,31 +1,26 @@
 from dcim.models import Region
-from ruamel.yaml import YAML
-from pathlib import Path
+from startup_script_utils import load_yaml
 import sys
 
-file = Path('/opt/netbox/initializers/regions.yml')
-if not file.is_file():
-  sys.exit()
+regions = load_yaml('/opt/netbox/initializers/regions.yml')
 
-with file.open('r') as stream:
-  yaml=YAML(typ='safe')
-  regions = yaml.load(stream)
+if regions is None:
+  sys.exit()
 
-  optional_assocs = {
-    'parent': (Region, 'name')
-  }
+optional_assocs = {
+  'parent': (Region, 'name')
+}
 
-  if regions is not None:
-    for params in regions:
+for params in regions:
 
-      for assoc, details in optional_assocs.items():
-        if assoc in params:
-          model, field = details
-          query = { field: params.pop(assoc) }
+  for assoc, details in optional_assocs.items():
+    if assoc in params:
+      model, field = details
+      query = { field: params.pop(assoc) }
 
-          params[assoc] = model.objects.get(**query)
+      params[assoc] = model.objects.get(**query)
 
-      region, created = Region.objects.get_or_create(**params)
+  region, created = Region.objects.get_or_create(**params)
 
-      if created:
-        print("🌐 Created region", region.name)
+  if created:
+    print("🌐 Created region", region.name)

startup_scripts/040_sites.py

@@ -1,46 +1,41 @@
 from dcim.models import Region, Site
 from extras.models import CustomField, CustomFieldValue
 from tenancy.models import Tenant
-from ruamel.yaml import YAML
-from pathlib import Path
+from startup_script_utils import load_yaml
 import sys
 
-file = Path('/opt/netbox/initializers/sites.yml')
-if not file.is_file():
-  sys.exit()
+sites = load_yaml('/opt/netbox/initializers/sites.yml')
 
-with file.open('r') as stream:
-  yaml = YAML(typ='safe')
-  sites = yaml.load(stream)
+if sites is None:
+  sys.exit()
 
-  optional_assocs = {
-    'region': (Region, 'name'),
-    'tenant': (Tenant, 'name')
-  }
+optional_assocs = {
+  'region': (Region, 'name'),
+  'tenant': (Tenant, 'name')
+}
 
-  if sites is not None:
-    for params in sites:
-      custom_fields = params.pop('custom_fields', None)
+for params in sites:
+  custom_fields = params.pop('custom_fields', None)
 
-      for assoc, details in optional_assocs.items():
-        if assoc in params:
-          model, field = details
-          query = { field: params.pop(assoc) }
+  for assoc, details in optional_assocs.items():
+    if assoc in params:
+      model, field = details
+      query = { field: params.pop(assoc) }
 
-          params[assoc] = model.objects.get(**query)
+      params[assoc] = model.objects.get(**query)
 
-      site, created = Site.objects.get_or_create(**params)
+  site, created = Site.objects.get_or_create(**params)
 
-      if created:
-        if custom_fields is not None:
-          for cf_name, cf_value in custom_fields.items():
-            custom_field = CustomField.objects.get(name=cf_name)
-            custom_field_value = CustomFieldValue.objects.create(
-              field=custom_field,
-              obj=site,
-              value=cf_value
-            )
+  if created:
+    if custom_fields is not None:
+      for cf_name, cf_value in custom_fields.items():
+        custom_field = CustomField.objects.get(name=cf_name)
+        custom_field_value = CustomFieldValue.objects.create(
+          field=custom_field,
+          obj=site,
+          value=cf_value
+        )
 
-            site.custom_field_values.add(custom_field_value)
+        site.custom_field_values.add(custom_field_value)
 
-        print("📍 Created site", site.name)
+    print("📍 Created site", site.name)

startup_scripts/050_manufacturers.py

@@ -1,19 +1,14 @@
 from dcim.models import Manufacturer
-from ruamel.yaml import YAML
-from pathlib import Path
+from startup_script_utils import load_yaml
 import sys
 
-file = Path('/opt/netbox/initializers/manufacturers.yml')
-if not file.is_file():
-  sys.exit()
+manufacturers = load_yaml('/opt/netbox/initializers/manufacturers.yml')
 
-with file.open('r') as stream:
-  yaml = YAML(typ='safe')
-  manufacturers = yaml.load(stream)
+if manufacturers is None:
+  sys.exit()
 
-  if manufacturers is not None:
-    for params in manufacturers:
-      manufacturer, created = Manufacturer.objects.get_or_create(**params)
+for params in manufacturers:
+  manufacturer, created = Manufacturer.objects.get_or_create(**params)
 
-      if created:
-        print("🏭 Created Manufacturer", manufacturer.name)
+  if created:
+    print("🏭 Created Manufacturer", manufacturer.name)