#!/usr/bin/env bash
set -euo pipefail

# root only
if [[ $EUID -ne 0 ]]; then
  echo "Must run as root."
  exit 1
fi

# install python/update system
if command -v apt-get &>/dev/null; then
  apt-get update
  apt-get install -y python3 python3-pip
elif command -v yum &>/dev/null; then
  yum update -y
  yum install -y python3 python3-pip
else
  echo "Unsupported distro."
  exit 1
fi

# install reqs
pip3 install --upgrade pip
pip3 install psutil distro flask prometheus_client

# create user + group
if ! id prometheus &>/dev/null; then
  useradd -r -s /sbin/nologin prometheus
fi

# create prometheus dir
mkdir -p /prometheus
chown prometheus:prometheus /prometheus

# prometheus endpoint (thanks, GPT)
cat > /prometheus/prometheus.py << 'EOF'
import os
import psutil
import platform
import distro
from functools import wraps
from flask import Flask, request, Response
from prometheus_client import CollectorRegistry, Gauge, Counter, generate_latest, CONTENT_TYPE_LATEST
from prometheus_client.core import GaugeMetricFamily, CounterMetricFamily

BASIC_AUTH_USER = os.environ.get('BASIC_AUTH_USER', 'admin')
BASIC_AUTH_PASSWORD = os.environ.get('BASIC_AUTH_PASSWORD', 'secret')

app = Flask(__name__)

def check_auth(username, password):
    return username == BASIC_AUTH_USER and password == BASIC_AUTH_PASSWORD

def authenticate():
    return Response(
        '401',
        401,
        {'WWW-Authenticate': 'Basic realm="Login Required"'}
    )

def requires_auth(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        auth = request.authorization
        if not auth or not check_auth(auth.username, auth.password):
            return authenticate()
        return f(*args, **kwargs)
    return decorated

class SystemMetricsCollector:
    """Custom Prometheus collector to gather system metrics."""
    def collect(self):
        kernel = platform.release()
        dist = f"{distro.name(pretty=True)} {distro.version(pretty=True)}".strip()
        info = GaugeMetricFamily(
            'node_system_info',
            'System information as labels',
            labels=['kernel_version', 'distribution']
        )
        info.add_metric([kernel, dist], 1)
        yield info

        cpu_percent = psutil.cpu_percent(interval=None)
        yield GaugeMetricFamily(
            'node_cpu_usage_percent',
            'Total CPU usage percent',
            value=cpu_percent
        )

        vm = psutil.virtual_memory()
        yield GaugeMetricFamily(
            'node_memory_usage_percent',
            'Total memory usage percent',
            value=vm.percent
        )

        net = psutil.net_io_counters()
        yield CounterMetricFamily(
            'node_network_bytes_sent_total',
            'Total network bytes sent',
            value=net.bytes_sent
        )
        yield CounterMetricFamily(
            'node_network_bytes_recv_total',
            'Total network bytes received',
            value=net.bytes_recv
        )

        proc_cpu = GaugeMetricFamily(
            'process_cpu_usage_percent',
            'Per-process CPU usage percent',
            labels=['pid', 'name']
        )
        proc_mem = GaugeMetricFamily(
            'process_memory_usage_percent',
            'Per-process memory usage percent',
            labels=['pid', 'name']
        )

        for proc in psutil.process_iter(['pid', 'name']):
            try:
                cpu = proc.cpu_percent(interval=None)
                mem = proc.memory_percent()
                pid = str(proc.info['pid'])
                name = proc.info['name'] or ''
                proc_cpu.add_metric([pid, name], cpu)
                proc_mem.add_metric([pid, name], mem)
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue

        yield proc_cpu
        yield proc_mem

# Register collector
registry = CollectorRegistry()
registry.register(SystemMetricsCollector())

@app.route('/metrics')
@requires_auth
def metrics():
    """Endpoint to expose Prometheus metrics with basic auth."""
    data = generate_latest(registry)
    return Response(data, mimetype=CONTENT_TYPE_LATEST)

if __name__ == '__main__':
    port = int(os.environ.get('METRICS_PORT', '9999'))
    app.run(host='127.0.0.1', port=port)
EOF

# 755
chown prometheus:prometheus /prometheus/prometheus.py
chmod 755 /prometheus/prometheus.py

# generate token
TOKEN=$(tr -dc 'A-Za-z0-9' </dev/urandom | head -c64)

# create service
cat > /etc/systemd/system/monitoring.service << EOF
[Unit]
Description=Monitoring Service - Prometheus Python Exporter
After=network.target

[Service]
Type=simple
WorkingDirectory=/prometheus
ExecStart=/usr/bin/python3 /prometheus/prometheus.py
Restart=on-failure
RestartSec=5s
Environment=BASIC_AUTH_USER=exporter BASIC_AUTH_PASSWORD=${TOKEN} METRICS_PORT=9999
User=prometheus
Group=prometheus

[Install]
WantedBy=multi-user.target
EOF

# services
systemctl daemon-reload
systemctl enable monitoring.service
systemctl start monitoring.service

echo "Monitoring setup complete."
echo "- Token: ${TOKEN}"
