Optimizing Code and Style

mdd/scripts/generate_inventory.py

@@ -1,88 +1,75 @@
 import os
 import requests
 import yaml
-
-def fetch_netbox_devices():
-    # Retrieve API key and URL from environment variables
-    netbox_api_url = os.environ.get('NETBOX_API')
-    netbox_api_token = os.environ.get('NETBOX_TOKEN')
-
-    if not netbox_api_url or not netbox_api_token:
-        print("Error: NETBOX_API or NETBOX_TOKEN environment variables are not set.")
-        exit(1)
-
-    # Construct the API request headers
-    headers = {
-        'Authorization': f'Token {netbox_api_token}'
+import load_env_vars
+
+ENV_VARS = {
+    "NAUTOBOT_URL" : None,
+    "NAUTOBOT_TOKEN" : None
+}
+
+ENV_VARS = load_env_vars.load_env_vars(os.environ, ENV_VARS)
+
+headers = {'Authorization': f'Token {ENV_VARS['NAUTOBOT_TOKEN']}'}
+verify_ssl = False
+existing_inventory_path = '/path/to/existing/inventory.yml'  # Replace with the actual path
+
+def get_inventory_structure(): #TODO: Compare to generate_tags
+    """Generated default inventory structure"""
+    return {
+        'all': {
+            'children': {
+                'network': {
+                    'children': {}
+                }
+            }
+        }
     }
 
-    # Disable SSL certificate validation
-    verify_ssl = False
+def fetch_netbox_devices():
+    existing_inventory = {}
 
     # Make API requests to NetBox to retrieve device information
     response = requests.get(
-        f'{netbox_api_url}/api/dcim/devices/',
+        f'{ENV_VARS['NAUTOBOT_URL']}/api/dcim/devices/',
         headers=headers,
         verify=verify_ssl  # Disable SSL certificate validation
     )
 
-    if response.status_code == 200:
-        # Parse the JSON response into a Python dictionary
-        device_data = response.json()
-
-        # Create an empty inventory dictionary
-        inventory_data = {
-            'all': {
-                'children': {
-                    'network': {
-                        'children': {}
-                    }
-                }
-            }
-        }
-
-        # Check for an existing inventory file in a different directory
-        existing_inventory_path = '/path/to/existing/inventory.yml'  # Replace with the actual path
-        existing_inventory = {}
-
-        if os.path.exists(existing_inventory_path):
-            with open(existing_inventory_path, 'r') as existing_inventory_file:
-                existing_inventory = yaml.safe_load(existing_inventory_file)
-
-        # Populate the inventory with devices based on device role
-        for device in device_data['results']:
-            device_name = device['name']
-            device_role = device['device_role']['slug']
-
-            # Check if the device is not in the existing inventory
-            if device_name not in existing_inventory.get('all', {}).get('children', {}).get('network', {}).get('children', {}).get(device_role, {}).get('hosts', {}):
-                # Check if primary_ip4 exists and is not None
-                primary_ip_data = device.get('primary_ip4')
-                if primary_ip_data and primary_ip_data.get('address'):
-                    primary_ip = primary_ip_data['address'].split('/')[0]  # Extract IP without the /32
-                    # Create groups for each device role under the "network" group if they don't exist
-                    if device_role not in inventory_data['all']['children']['network']['children']:
-                        inventory_data['all']['children']['network']['children'][device_role] = {
-                            'hosts': {}
-                        }
-
-                    # Add the device to its respective role group and include the ansible_host
-                    inventory_data['all']['children']['network']['children'][device_role]['hosts'][device_name] = {
-                        'ansible_host': primary_ip
-                    }
-                else:
-                    # If there is no IP address available, create groups without ansible_host
-                    if device_role not in inventory_data['all']['children']['network']['children']:
-                        inventory_data['all']['children']['network']['children'][device_role] = {
-                            'hosts': {}
-                        }
-                    inventory_data['all']['children']['network']['children'][device_role]['hosts'][device_name] = {}
-
-        return inventory_data
-    else:
-        print(f"Error fetching data from NetBox: {response.status_code}")
+    if response.status_code != 200:
+        print(f"Error fetching devices from NetBox: {response.status_code}")
         exit(1)
 
+    # Create an empty inventory dictionary
+    inventory_data = get_inventory_structure()
+
+    # Check for an existing inventory file in a different directory
+    if os.path.exists(existing_inventory_path):
+        with open(existing_inventory_path, 'r') as existing_inventory_file:
+            existing_inventory = yaml.safe_load(existing_inventory_file)
+
+    # Populate the inventory with devices based on device role
+    for device in response.json()['results']:
+        device_name = device['name']
+        device_role = device['device_role']['slug']
+
+        # Check if the device is not in the existing inventory
+        if device_name in existing_inventory.get('all', {}).get('children', {}).get('network', {}).get('children', {}).get(device_role, {}).get('hosts', {}):
+            continue
+
+        # Check if primary_ip4 exists and is not None
+        primary_ip_data = device.get('primary_ip4')
+        if not primary_ip_data: # Still need to create the role so that it shows up in the inventory file - because line 69 won't hit these cases
+            inventory_data['all']['children']['network']['children'].setdefault(device_role, {}).setdefault('hosts', {}).setdefault(device_name, {})
+            continue
+
+        primary_ip_address = primary_ip_data.get('address') #TODO: will this ever not be there
+        primary_ip = primary_ip_address.split('/')[0]  # Extract IP without the /32
+
+        inventory_data['all']['children']['network']['children'].setdefault(device_role, {}).setdefault('hosts', {}).setdefault(device_name, {})['ansible_host'] = primary_ip
+
+    return inventory_data
+
 def write_inventory_file(inventory_data):
     # Define the inventory file path as "inventory/network.yml"
     inventory_file_path = os.path.join(os.path.dirname(__file__), '..', 'inventory', 'network.yml')

mdd/scripts/generate_tags.py

@@ -1,21 +1,22 @@
 import os
 import requests
 import yaml
-import json
+import load_env_vars
 
-def fetch_netbox_devices():
-    netbox_api_url = os.environ.get('NETBOX_API')
-    netbox_api_token = os.environ.get('NETBOX_TOKEN')
+ENV_VARS = {
+    "NAUTOBOT_URL" : None,
+    "NAUTOBOT_TOKEN" : None
+}
 
-    if not netbox_api_url or not netbox_api_token:
-        print("NETBOX_API or NETBOX_TOKEN environment variables are not set.")
-        exit(1)
+ENV_VARS = load_env_vars.load_env_vars(os.environ, ENV_VARS)
 
-    headers = {'Authorization': f'Token {netbox_api_token}', 'Accept': 'application/json'}
-    verify_ssl = False
+headers = {'Authorization': f'Token {ENV_VARS['NAUTOBOT_TOKEN']}', 'Accept': 'application/json'}
+verify_ssl = False
 
-    # Initialize the inventory structure with vars for sites
-    inventory_data = {
+def get_inventory_structure():
+    """Returns the default inventory structure"""
+
+    return {
         'all': {
             'vars': {
                 'sites': []
@@ -28,36 +29,40 @@ def fetch_netbox_devices():
         }
     }
 
+def fetch_netbox_devices():
+    # Initialize the inventory structure with vars for sites
+    inventory_data = get_inventory_structure()
+
+    # Fetch sites from NetBox
+    response = requests.get(f'{ENV_VARS['NAUTOBOT_URL']}/api/dcim/sites/?limit=1000', headers=headers, verify=verify_ssl)
+
+    if response.status_code != 200:
+        print(f"Failed to fetch sites from NetBox: {response.status_code}")
+        exit(1)
+
+    # Populate the sites list
+    for site in response.json()['results']:
+        inventory_data['all']['vars']['sites'].append(site['name'])
+
     # Fetch devices from NetBox
-    response = requests.get(f'{netbox_api_url}/api/dcim/sites/?limit=1000', headers=headers, verify=verify_ssl)
-    if response.status_code == 200:
-        sites_data = response.json()
-        # Populate the sites list
-        for site in sites_data['results']:
-            inventory_data['all']['vars']['sites'].append(site['name'])
-
-    response = requests.get(f'{netbox_api_url}/api/dcim/devices/?limit=1000', headers=headers, verify=verify_ssl)
-    if response.status_code == 200:
-        device_data = response.json()
-
-        for device in device_data['results']:
-            device_name = device['name']
-            site_name = device.get('site', {}).get('name', 'undefined')
-            device_tags = [tag['name'] for tag in device.get('tags', [])]
-
-            if site_name not in inventory_data['all']['children']['org']['children']:
-                inventory_data['all']['children']['org']['children'][site_name] = {'hosts': {}}
-
-            # Add the device under the site with its tags
-            inventory_data['all']['children']['org']['children'][site_name]['hosts'][device_name] = {'tags': device_tags}
-    else:
-        print(f"Failed to fetch data from NetBox: {response.status_code}")
+    response = requests.get(f'{ENV_VARS['NAUTOBOT_URL']}/api/dcim/devices/?limit=1000', headers=headers, verify=verify_ssl)
+
+    if response.status_code != 200:
+        print(f"Failed to fetch devices from NetBox: {response.status_code}")
         exit(1)
 
+    for device in response.json()['results']:
+        device_name = device['name']
+        site_name = device.get('site', {}).get('name', 'undefined')
+        device_tags = [tag['name'] for tag in device.get('tags', [])]
+
+        # Add the device under the site with its tags
+        inventory_data['all']['children']['org']['children'].setdefault(site_name, {}).setdefault('hosts', {}).setdefault(device_name, {})['tags'] = device_tags
+
     return inventory_data
 
 def write_inventory_file(inventory_data):
-    inventory_file_path = os.path.join(os.path.dirname(__file__),'..', 'inventory', 'tags.yml')
+    inventory_file_path = os.path.join(os.path.dirname(__file__), '..', 'inventory', 'tags.yml')
 
     with open(inventory_file_path, 'w') as f:
         yaml.dump(inventory_data, f, default_flow_style=False, sort_keys=False)

mdd/scripts/load_env_vars.py

@@ -0,0 +1,33 @@
+# CALL FILE TOOLS?
+def load_env_vars(OS_ENV, ENV_VARS):
+    """
+    Sees if any environment variables are not set - if they are not, exits the program
+    OS_ENV should be passed as os.environ
+    """
+    missing_vars = []
+    for var in ENV_VARS:
+        value = OS_ENV.get(var)
+        if not value:
+            missing_vars.append(var)
+        else:
+            # Create a variable with the name from the environment variable
+            ENV_VARS[var] = value
+
+    # If any required environment variable is missing or empty, print an error message and exit
+    if missing_vars:
+        print(f"Error: The following environment variables are not set or empty: {', '.join(missing_vars)}")
+        exit(1)
+
+    return ENV_VARS
+
+def get_sub_dict(dict, level):
+    """
+    Returns reference to a sub dictionary based on first key
+    Used for inventory where you only have 1 key
+    """
+    sub_dict = dict
+    for _ in range(level):
+        key = list(sub_dict.keys())[0] # get first key
+        sub_dict = sub_dict[key]
+
+    return sub_dict # reference