Refactor Prometheus config update function and improve authentication security

Author: codeperfectplus

src/routes/helper/prometheus_helper.py

@@ -1,6 +1,7 @@
 import os
 import yaml
 import subprocess
+from collections import OrderedDict
 from src.utils import ROOT_DIR
 
 prometheus_yml_path = os.path.join(ROOT_DIR, 'prometheus_config/prometheus.yml')
@@ -16,82 +17,103 @@ def is_valid_file(file_path: str) -> bool:
                 return False
     return True
 
+
+class OrderedDumper(yaml.SafeDumper):
+    """Custom YAML dumper that preserves order of keys."""
+    pass
+
+def dict_representer(dumper, data):
+    return dumper.represent_dict(data.items())
+
+OrderedDumper.add_representer(OrderedDict, dict_representer)
+
 def load_yaml(file_path):
-    """Load the YAML file."""
+    """Load the YAML file and keep the order of dictionaries."""
     with open(file_path, 'r') as file:
-        return yaml.safe_load(file)
+        return yaml.load(file, Loader=yaml.SafeLoader)
 
 def save_yaml(data, file_path):
-    """Save the updated YAML data back to the file."""
+    """Save the updated YAML data back to the file, preserving order."""
     with open(file_path, 'w') as file:
-        yaml.dump(data, file, default_flow_style=False)
-
-def show_targets():
-    """Show all targets for each job."""
-    config = load_yaml(prometheus_yml_path)
-    targets_info = []
-    for scrape_config in config.get('scrape_configs', []):
-        job_name = scrape_config['job_name']
-        targets = scrape_config['static_configs'][0]['targets']
-        scrape_interval = scrape_config.get('scrape_interval', '15s')
-        targets_info.append({
-            'job_name': job_name,
-            'targets': targets,
-            'scrape_interval': scrape_interval
-        })
-    return targets_info
+        yaml.dump(data, file, Dumper=OrderedDumper, default_flow_style=False)
 
 def update_prometheus_config():
-    """
-    Update the first target whenever the function is called.
-    If the network changes, the IP address will also change.
-    This updates the first IP address in the list, assuming it's the machine with the 'localhost' job.
-    """
+    """Update the first target with the machine's IP address."""
     print("Updating Prometheus config...")
 
     # Get the machine's IP address
-    ipv4_address = subprocess.run(['hostname', '-I'], capture_output=True, text=True, check=True).stdout.split()[0]
-    
-    # Define the path to Prometheus config file
-    prometheus_config_path = os.path.join(ROOT_DIR, 'prometheus_config', 'prometheus.yml')
-    
+    try:
+        ipv4_address = subprocess.run(['hostname', '-I'], capture_output=True, text=True, check=True).stdout.split()[0]
+    except subprocess.CalledProcessError as e:
+        print(f"Error getting IP address: {e}")
+        return False
+
     # Load the existing config
-    config = load_yaml(prometheus_config_path)
+    try:
+        config = load_yaml(prometheus_yml_path)
+    except Exception as e:
+        print(f"Error loading YAML config: {e}")
+        return False
+
     # Fetch the 'localhost' job
-    localhost_job = next((job for job in config['scrape_configs'] if job['job_name'] == 'localhost'), None)
+    localhost_job = next((job for job in config.get('scrape_configs', []) if job.get('job_name') == 'localhost'), None)
 
     if localhost_job:
         # Update the IP address for the 'localhost' job target
         localhost_job['static_configs'][0]['targets'][0] = f'{ipv4_address}:5050'
 
-        # localhost_job['basic_auth'] = {
-        #     'username': "username",
-        #     'password': "password"
-        # }
-        
-        # Save the updated config
-        save_yaml(config, prometheus_config_path)
+        # Create a new OrderedDict to maintain the correct order
+        updated_job = OrderedDict()
+        updated_job['job_name'] = localhost_job['job_name']
+        updated_job['scrape_interval'] = localhost_job.get('scrape_interval', '10s')
+        updated_job['static_configs'] = localhost_job['static_configs']
 
-        print("Prometheus config updated successfully.")
-        return True
+        # Add basic_auth last to maintain order
+        if 'basic_auth' in localhost_job:
+            updated_job['basic_auth'] = localhost_job['basic_auth']
+
+        # Replace the old job with the updated one
+        for index, job in enumerate(config['scrape_configs']):
+            if job['job_name'] == 'localhost':
+                config['scrape_configs'][index] = updated_job
+                break
 
+        # Save the updated config
+        try:
+            save_yaml(config, prometheus_yml_path)
+            print("Prometheus config updated successfully.")
+            return True
+        except Exception as e:
+            print(f"Error saving YAML config: {e}")
+            return False
+    
     print("No 'localhost' job found in Prometheus config.")
     return False
 
+def show_targets():
+    """Show all targets for each job."""
+    config = load_yaml(prometheus_yml_path)
+    targets_info = []
+    for scrape_config in config.get('scrape_configs', []):
+        job_name = scrape_config['job_name']
+        targets = scrape_config.get('static_configs', [{}])[0].get('targets', [])
+        scrape_interval = scrape_config.get('scrape_interval', '15s')
+        targets_info.append({
+            'job_name': job_name,
+            'targets': targets,
+            'scrape_interval': scrape_interval
+        })
+    return targets_info
+
 def update_prometheus_container():
     """Update the Prometheus container."""
-    # Define the path to your shell script
     try:
-        # Use subprocess.run to execute the shell script
         result = subprocess.run(['bash', update_prometheus_path], check=True, text=True, capture_output=True)
-
-        # Print the output of the script
         print("Output:")
         print(result.stdout)
 
-        # Print any errors (if any)
         if result.stderr:
             print("Errors:")
             print(result.stderr)
     except subprocess.CalledProcessError as e:
-        print(f"An error occurred: {e}")
+        print(f"An error occurred while updating Prometheus container: {e}")

src/routes/prometheus.py

@@ -2,6 +2,7 @@
 from prometheus_client import generate_latest
 import os
 import yaml
+from collections import OrderedDict
 
 from src.config import app, db
 from src.models import ExternalMonitornig
@@ -73,14 +74,13 @@ def delete_file_path(id):
 @app.route('/configure_targets')
 @admin_required
 def configure_targets():
-    # update_prometheus_config()
+    update_prometheus_config()
     targets_info = show_targets()
     return render_template('other/targets.html', targets_info=targets_info)
 
 @app.route('/targets/restart_prometheus')
 @admin_required
 def restart_prometheus():
-    update_prometheus_config
     update_prometheus_container()
     flash('Prometheus container restarted successfully!', 'success')
     return redirect(url_for('configure_targets'))
@@ -89,26 +89,63 @@ def restart_prometheus():
 def add_target():
     job_name = request.form.get('job_name')
     new_target = request.form.get('new_target')
+    username = request.form.get('username')
+    password = request.form.get('password')
     scrape_interval = request.form.get('scrape_interval', '15s') + 's'  # New scrape interval
     config = load_yaml(prometheus_yml_path)
 
-    # new target should be <ip>:<port> check if it is in the correct format
+    # Validate target format
     if ':' not in new_target:
         flash('Invalid target format. It should be in the format <ip>:<port>.', 'danger')
         return redirect(url_for('configure_targets'))
-    
+
+    job_found = False
+
     for scrape_config in config['scrape_configs']:
         if scrape_config['job_name'] == job_name:
+            # Append new target
             scrape_config['static_configs'][0]['targets'].append(new_target)
+            job_found = True
+            
+            # Update scrape interval
+            scrape_config['scrape_interval'] = scrape_interval
+            
+            # Prepare the updated job dictionary to maintain order
+            updated_job = OrderedDict()
+            updated_job['job_name'] = scrape_config['job_name']
+            updated_job['static_configs'] = scrape_config['static_configs']
+            updated_job['scrape_interval'] = scrape_config['scrape_interval']
+
+            # Update basic_auth if provided
+            if username and password:
+                updated_job['basic_auth'] = {
+                    'username': username,
+                    'password': password
+                }
+
+            # Replace the existing job with the updated one
+            index = config['scrape_configs'].index(scrape_config)
+            config['scrape_configs'][index] = updated_job
+            
             break
-    else:
-        new_job = {
-            'job_name': job_name,
-            'static_configs': [{'targets': [new_target]}],
-            'scrape_interval': scrape_interval  # Set the specific interval
-        }
-        config['scrape_configs'].append(new_job)
+
+    # if not job_found:
+    #     # Create new job entry
+    #     new_job = OrderedDict()
+    #     new_job['job_name'] = job_name
+    #     new_job['static_configs'] = [{'targets': [new_target]}]
+    #     new_job['scrape_interval'] = scrape_interval
 
+    #     # Add basic_auth if provided
+    #     if username and password:
+    #         new_job['basic_auth'] = {
+    #             'username': username,
+    #             'password': password
+    #         }
+    #     # Append the new job to scrape_configs
+    #     config['scrape_configs'].append(new_job)
+
+    # Save the updated config
     save_yaml(config, prometheus_yml_path)
     flash('Target added successfully!', 'success')
     # update_prometheus_container()

src/scripts/prometheus.sh

@@ -27,7 +27,7 @@ PROMETHEUS_DATA_DIR="/home/$USER_NAME/.database/prometheus"
 FLASK_APP_IP=$(hostname -I | cut -d' ' -f1)
 FLASK_APP_PORT="5050"
 SCRAPING_INTERVAL="10s"
-JOB_NAME='systemguard-metrics'
+monitor='systemguard-metrics'
 
 # Logging function for better readability
 log() {
@@ -42,6 +42,12 @@ mkdir -p "$PROMETHEUS_DATA_DIR"
 # Create the prometheus.yml configuration file
 log "Generating prometheus.yml configuration file."
 cat > "$PROMETHEUS_CONFIG_FILE" <<EOL
+global:
+  external_labels:
+    system: $monitor
+    environment: 'production'
+    user: '$USER_NAME'
+    
 scrape_configs:
   - job_name: localhost
     scrape_interval: $SCRAPING_INTERVAL