#!/bin/bash
################################################################################
# Script Name: add-http-auth.sh
# Version: 3.0
# Description: Add HTTP Basic Auth to Prometheus stack reverse proxies
#              Supports both nginx and Apache — auto-detects which is in use.
#              Uses non-destructive include snippets to preserve existing
#              HTTPS/certbot configs.
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Supported Services:
#   - Prometheus (port 9090)
#   - Alertmanager (port 9093)
#   - Mimir (port 9009) — optionally protects /api/v1/push
#   - Loki (port 3100) — optionally protects /loki/api/v1/push
#
# Supported Web Servers:
#   - nginx  — inserts 'include' snippets into location blocks
#   - Apache — inserts 'Include' snippets into <Location> blocks
#
# Usage:
#   sudo ./add-http-auth.sh
#   sudo ./add-http-auth.sh --remove
#   sudo ./add-http-auth.sh --status
#
################################################################################

set -euo pipefail

SCRIPT_VERSION="3.0"
BACKUP_DIR="/var/backups/http-auth"

# Detected at runtime
WEB_SERVER=""       # "nginx" or "apache"
CONFIG_DIR=""       # where vhost configs live
SNIPPET_DIR=""      # where auth snippets go
AUTH_DIR=""         # where htpasswd files go
WEB_USER=""         # www-data, nginx, apache, etc.
SERVICE_NAME=""     # systemd service name

# Service definitions: name|nginx_config|apache_config|port
SERVICES=(
    "prometheus|prometheus.conf|prometheus.conf|9090"
    "alertmanager|alerts.conf|alerts.conf|9093"
    "mimir|mimir.conf|mimir.conf|9009"
    "loki|loki.conf|loki.conf|3100"
)

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

show_usage() {
    cat <<EOF
Usage: $0 [OPTIONS]

Add HTTP Basic Auth to Prometheus stack reverse proxies (v${SCRIPT_VERSION}).
Supports both nginx and Apache — auto-detects which is in use.

OPTIONS:
    --remove    Remove auth from all services (restore backups)
    --status    Show current auth status for each service
    -h, --help  Show this help

EXAMPLES:
    $0              # Interactive setup
    $0 --status     # Check which services have auth enabled
    $0 --remove     # Remove auth and restore original configs

EOF
    exit 0
}

die() {
    echo "ERROR: $1" >&2
    exit 1
}

warn() {
    echo "WARNING: $1" >&2
}

# Get the config filename for the current web server
get_config_file() {
    local entry="$1"
    local name nginx_conf apache_conf port
    IFS='|' read -r name nginx_conf apache_conf port <<< "$entry"
    if [ "$WEB_SERVER" = "nginx" ]; then
        echo "$nginx_conf"
    else
        echo "$apache_conf"
    fi
}

get_service_name() {
    local entry="$1"
    IFS='|' read -r name _ _ _ <<< "$entry"
    echo "$name"
}

get_service_port() {
    local entry="$1"
    IFS='|' read -r _ _ _ port <<< "$entry"
    echo "$port"
}

# ============================================================================
# WEB SERVER DETECTION
# ============================================================================

detect_web_server() {
    local has_nginx=false
    local has_apache=false

    if command -v nginx &>/dev/null && systemctl is-active --quiet nginx 2>/dev/null; then
        has_nginx=true
    fi

    if command -v apache2ctl &>/dev/null && systemctl is-active --quiet apache2 2>/dev/null; then
        has_apache=true
    elif command -v httpd &>/dev/null && systemctl is-active --quiet httpd 2>/dev/null; then
        has_apache=true
    fi

    if [ "$has_nginx" = true ] && [ "$has_apache" = true ]; then
        echo ""
        echo "Both nginx and Apache detected. Which are you using for reverse proxies?"
        echo "  1) nginx"
        echo "  2) Apache"
        read -r -p "Select [1]: " choice
        case "${choice:-1}" in
            2) WEB_SERVER="apache" ;;
            *) WEB_SERVER="nginx" ;;
        esac
    elif [ "$has_nginx" = true ]; then
        WEB_SERVER="nginx"
    elif [ "$has_apache" = true ]; then
        WEB_SERVER="apache"
    else
        die "Neither nginx nor Apache detected as running"
    fi

    echo "  Detected web server: ${WEB_SERVER}"
}

# Set paths based on detected web server
configure_paths() {
    if [ "$WEB_SERVER" = "nginx" ]; then
        if [ -d "/etc/nginx/sites-available" ]; then
            CONFIG_DIR="/etc/nginx/sites-available"
        elif [ -d "/etc/nginx/conf.d" ]; then
            CONFIG_DIR="/etc/nginx/conf.d"
        else
            die "nginx config directory not found"
        fi
        SNIPPET_DIR="/etc/nginx/snippets"
        AUTH_DIR="/etc/nginx/auth"
        SERVICE_NAME="nginx"

        if id "www-data" &>/dev/null; then
            WEB_USER="www-data"
        elif id "nginx" &>/dev/null; then
            WEB_USER="nginx"
        else
            WEB_USER="root"
        fi
    else
        # Apache
        if [ -d "/etc/apache2/sites-available" ]; then
            CONFIG_DIR="/etc/apache2/sites-available"
            SNIPPET_DIR="/etc/apache2/conf-available"
            SERVICE_NAME="apache2"
        elif [ -d "/etc/httpd/conf.d" ]; then
            CONFIG_DIR="/etc/httpd/conf.d"
            SNIPPET_DIR="/etc/httpd/conf.d"
            SERVICE_NAME="httpd"
        else
            die "Apache config directory not found"
        fi
        AUTH_DIR="/etc/httpd/auth"
        [ -d "/etc/apache2" ] && AUTH_DIR="/etc/apache2/auth"

        if id "www-data" &>/dev/null; then
            WEB_USER="www-data"
        elif id "apache" &>/dev/null; then
            WEB_USER="apache"
        else
            WEB_USER="root"
        fi
    fi
}

# ============================================================================
# HTTPS DETECTION
# ============================================================================

has_https() {
    local config_file="$1"
    if [ "$WEB_SERVER" = "nginx" ]; then
        grep -qE 'listen\s+.*443\s+ssl' "$config_file" 2>/dev/null
    else
        grep -qE 'SSLEngine\s+on|<VirtualHost\s+\*:443' "$config_file" 2>/dev/null
    fi
}

# ============================================================================
# AUTH SNIPPET CHECK
# ============================================================================

has_auth_snippet() {
    local config_file="$1"
    local service="$2"
    if [ "$WEB_SERVER" = "nginx" ]; then
        grep -qF "include ${SNIPPET_DIR}/auth-${service}.conf" "$config_file" 2>/dev/null
    else
        grep -qF "Include ${SNIPPET_DIR}/auth-${service}.conf" "$config_file" 2>/dev/null
    fi
}

# ============================================================================
# SETUP FUNCTIONS
# ============================================================================

install_htpasswd() {
    if command -v htpasswd &>/dev/null; then
        return 0
    fi

    echo "Installing htpasswd..."
    if command -v apt-get &>/dev/null; then
        apt-get -y install apache2-utils
    elif command -v dnf &>/dev/null; then
        dnf -y install httpd-tools
    elif command -v yum &>/dev/null; then
        yum -y install httpd-tools
    else
        die "Cannot install htpasswd — install apache2-utils or httpd-tools manually"
    fi
}

backup_config() {
    local config_file="$1"
    local timestamp
    timestamp=$(date +%F_%H%M%S)
    local backup_path="${BACKUP_DIR}/${timestamp}"

    mkdir -p "$backup_path"
    cp "$config_file" "$backup_path/"
    echo "  Backed up to ${backup_path}/$(basename "$config_file")"
}

# ============================================================================
# NGINX-SPECIFIC FUNCTIONS
# ============================================================================

nginx_create_snippet() {
    local service="$1"
    local display_name="$2"

    cat > "${SNIPPET_DIR}/auth-${service}.conf" <<EOF
# Auth snippet for ${display_name} — managed by add-http-auth.sh
auth_basic "${display_name} - Authentication Required";
auth_basic_user_file ${AUTH_DIR}/.htpasswd-${service};
EOF

    echo "  Created ${SNIPPET_DIR}/auth-${service}.conf"
}

nginx_insert_auth() {
    local config_file="$1"
    local service="$2"
    local include_line="        include ${SNIPPET_DIR}/auth-${service}.conf;"

    if has_auth_snippet "$config_file" "$service"; then
        echo "  Auth already configured in $(basename "$config_file") — skipping"
        return 0
    fi

    backup_config "$config_file"

    local temp_file
    temp_file=$(mktemp)

    awk -v inc="$include_line" '
        /location \/ \{/ && !done {
            print
            print inc
            done = 1
            next
        }
        { print }
    ' "$config_file" > "$temp_file"

    mv "$temp_file" "$config_file"
    echo "  Inserted auth include into $(basename "$config_file")"
}

nginx_insert_push_auth() {
    local config_file="$1"
    local service="$2"

    if grep -q "location.*/api/v1/push" "$config_file" && \
       ! grep -A2 "location.*/api/v1/push" "$config_file" | grep -qF "auth-${service}.conf"; then
        local temp_file
        temp_file=$(mktemp)
        local include_line="        include ${SNIPPET_DIR}/auth-${service}.conf;"
        awk -v inc="$include_line" '
            /location.*\/api\/v1\/push/ && !push_done {
                print
                print inc
                push_done = 1
                next
            }
            { print }
        ' "$config_file" > "$temp_file"
        mv "$temp_file" "$config_file"
        echo "  Protected push endpoint with auth"
    fi
}

nginx_remove_auth() {
    local config_file="$1"
    local service="$2"

    local temp_file
    temp_file=$(mktemp)
    grep -vF "include ${SNIPPET_DIR}/auth-${service}.conf" "$config_file" > "$temp_file"
    mv "$temp_file" "$config_file"
}

nginx_test_config() {
    nginx -t 2>&1
}

# ============================================================================
# APACHE-SPECIFIC FUNCTIONS
# ============================================================================

apache_create_snippet() {
    local service="$1"
    local display_name="$2"

    cat > "${SNIPPET_DIR}/auth-${service}.conf" <<EOF
# Auth snippet for ${display_name} — managed by add-http-auth.sh
AuthType Basic
AuthName "${display_name} - Authentication Required"
AuthUserFile ${AUTH_DIR}/.htpasswd-${service}
Require valid-user
EOF

    echo "  Created ${SNIPPET_DIR}/auth-${service}.conf"
}

apache_insert_auth() {
    local config_file="$1"
    local service="$2"
    local include_line="    Include ${SNIPPET_DIR}/auth-${service}.conf"

    if has_auth_snippet "$config_file" "$service"; then
        echo "  Auth already configured in $(basename "$config_file") — skipping"
        return 0
    fi

    backup_config "$config_file"

    # Check if config uses <Location /> or <Proxy *>
    local temp_file
    temp_file=$(mktemp)

    if grep -qE '<Location\s+/\s*>' "$config_file"; then
        # Insert after <Location /> opening tag
        awk -v inc="$include_line" '
            /<Location\s+\/\s*>/ && !done {
                print
                print inc
                done = 1
                next
            }
            { print }
        ' "$config_file" > "$temp_file"
    elif grep -qE '<Proxy\s+' "$config_file"; then
        # Insert inside <VirtualHost> before the first ProxyPass
        awk -v inc="$include_line" -v sdir="${SNIPPET_DIR}" -v svc="$service" '
            /ProxyPass\s/ && !done {
                # Add a Location block with auth before ProxyPass
                print "    <Location />"
                print inc
                print "    </Location>"
                print ""
                done = 1
            }
            { print }
        ' "$config_file" > "$temp_file"
    else
        # No Location or Proxy block found — add a Location block before </VirtualHost>
        awk -v inc="$include_line" '
            /<\/VirtualHost>/ && !done {
                print ""
                print "    <Location />"
                print inc
                print "    </Location>"
                print ""
                done = 1
            }
            { print }
        ' "$config_file" > "$temp_file"
    fi

    mv "$temp_file" "$config_file"
    echo "  Inserted auth into $(basename "$config_file")"
}

apache_insert_push_auth() {
    local config_file="$1"
    local service="$2"
    local push_path=""

    if [ "$service" = "mimir" ]; then
        push_path="/api/v1/push"
    elif [ "$service" = "loki" ]; then
        push_path="/loki/api/v1/push"
    else
        return 0
    fi

    # Check if there's already a Location block for the push path
    if grep -qF "$push_path" "$config_file" && \
       ! grep -A3 "$push_path" "$config_file" | grep -qF "auth-${service}.conf"; then
        backup_config "$config_file"
        local temp_file
        temp_file=$(mktemp)
        local include_line="    Include ${SNIPPET_DIR}/auth-${service}.conf"
        awk -v path="$push_path" -v inc="$include_line" '
            $0 ~ path && /Location/ && !push_done {
                print
                print inc
                push_done = 1
                next
            }
            { print }
        ' "$config_file" > "$temp_file"
        mv "$temp_file" "$config_file"
        echo "  Protected push endpoint with auth"
    fi
}

apache_remove_auth() {
    local config_file="$1"
    local service="$2"

    local temp_file
    temp_file=$(mktemp)
    grep -vF "Include ${SNIPPET_DIR}/auth-${service}.conf" "$config_file" > "$temp_file"
    mv "$temp_file" "$config_file"
}

apache_test_config() {
    if command -v apache2ctl &>/dev/null; then
        apache2ctl configtest 2>&1
    else
        httpd -t 2>&1
    fi
}

# ============================================================================
# GENERIC WRAPPERS (dispatch to nginx or apache functions)
# ============================================================================

create_snippet() {
    if [ "$WEB_SERVER" = "nginx" ]; then
        nginx_create_snippet "$@"
    else
        apache_create_snippet "$@"
    fi
}

insert_auth() {
    if [ "$WEB_SERVER" = "nginx" ]; then
        nginx_insert_auth "$@"
    else
        apache_insert_auth "$@"
    fi
}

insert_push_auth() {
    if [ "$WEB_SERVER" = "nginx" ]; then
        nginx_insert_push_auth "$@"
    else
        apache_insert_push_auth "$@"
    fi
}

remove_auth_from_config() {
    if [ "$WEB_SERVER" = "nginx" ]; then
        nginx_remove_auth "$@"
    else
        apache_remove_auth "$@"
    fi
}

test_config() {
    if [ "$WEB_SERVER" = "nginx" ]; then
        nginx_test_config
    else
        apache_test_config
    fi
}

# ============================================================================
# STATUS & REMOVE
# ============================================================================

show_status() {
    detect_web_server
    configure_paths

    echo ""
    echo "=========================================="
    echo "HTTP Basic Auth Status (${WEB_SERVER})"
    echo "=========================================="
    echo ""

    for entry in "${SERVICES[@]}"; do
        local name config_file
        name=$(get_service_name "$entry")
        config_file=$(get_config_file "$entry")
        local display_name
        display_name="${name^}"
        local full_path="${CONFIG_DIR}/${config_file}"

        printf "  %-14s " "${display_name}:"

        if [ ! -f "$full_path" ]; then
            echo "no config found"
            continue
        fi

        if has_auth_snippet "$full_path" "$name"; then
            if [ -f "${AUTH_DIR}/.htpasswd-${name}" ]; then
                echo "ENABLED (htpasswd + snippet)"
            else
                echo "BROKEN (snippet exists but htpasswd file missing)"
            fi
        else
            echo "not configured"
        fi
    done

    echo ""
    echo "Web server:     ${WEB_SERVER}"
    echo "Config dir:     ${CONFIG_DIR}"
    echo "Snippet dir:    ${SNIPPET_DIR}"
    echo "Auth dir:       ${AUTH_DIR}"
    echo "Backup dir:     ${BACKUP_DIR}"
    echo ""
}

do_remove() {
    detect_web_server
    configure_paths

    echo ""
    echo "Removing HTTP Basic Auth from all services (${WEB_SERVER})..."
    echo ""

    for entry in "${SERVICES[@]}"; do
        local name config_file
        name=$(get_service_name "$entry")
        config_file=$(get_config_file "$entry")
        local full_path="${CONFIG_DIR}/${config_file}"

        if [ ! -f "$full_path" ]; then
            continue
        fi

        if has_auth_snippet "$full_path" "$name"; then
            backup_config "$full_path"
            remove_auth_from_config "$full_path" "$name"
            echo "  Removed auth from ${config_file}"
        fi

        rm -f "${SNIPPET_DIR}/auth-${name}.conf"
    done

    echo ""
    echo "Testing ${WEB_SERVER} configuration..."
    if test_config; then
        systemctl reload "$SERVICE_NAME"
        echo ""
        echo "Auth removed and ${WEB_SERVER} reloaded."
    else
        warn "${WEB_SERVER} config test failed — check your config manually"
    fi
}

# ============================================================================
# MAIN SETUP
# ============================================================================

setup_auth() {
    detect_web_server
    configure_paths

    echo ""
    echo "=========================================="
    echo "Add HTTP Basic Auth to Prometheus Stack"
    echo "Version: ${SCRIPT_VERSION} (${WEB_SERVER})"
    echo "=========================================="

    # Check for HTTPS
    local has_any_https=false
    for entry in "${SERVICES[@]}"; do
        local name config_file
        name=$(get_service_name "$entry")
        config_file=$(get_config_file "$entry")
        local full_path="${CONFIG_DIR}/${config_file}"
        if [ -f "$full_path" ] && has_https "$full_path"; then
            has_any_https=true
            break
        fi
    done

    if [ "$has_any_https" = false ]; then
        echo ""
        warn "No HTTPS configuration detected!"
        echo "  Basic Auth over HTTP sends credentials in cleartext."
        echo "  Strongly recommended: run certbot first to enable HTTPS."
        echo ""
        read -r -p "Continue without HTTPS? [y/N]: " confirm
        if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
            echo "Aborted. Run certbot first, then re-run this script."
            exit 0
        fi
    fi

    # Detect which services have configs
    echo ""
    echo "Detected services:"
    local found_any=false
    for entry in "${SERVICES[@]}"; do
        local name config_file
        name=$(get_service_name "$entry")
        config_file=$(get_config_file "$entry")
        local full_path="${CONFIG_DIR}/${config_file}"
        if [ -f "$full_path" ]; then
            local https_status="HTTP"
            has_https "$full_path" && https_status="HTTPS"
            echo "  ✓ ${name} (${config_file}) [${https_status}]"
            found_any=true
        fi
    done

    if [ "$found_any" = false ]; then
        die "No service configs found in ${CONFIG_DIR}. Set up ${WEB_SERVER} reverse proxies first."
    fi

    echo ""

    # Ask about push endpoint protection
    local protect_push=false
    echo "Mimir and Loki have push endpoints used by remote agents."
    echo "Protecting them requires configuring credentials in Prometheus/Alloy."
    read -r -p "Protect push endpoints with auth too? [y/N]: " push_confirm
    if [[ "$push_confirm" =~ ^[Yy]$ ]]; then
        protect_push=true
    fi

    # Ask about shared vs per-service credentials
    local shared_creds=false
    local shared_htpasswd=""
    echo ""
    echo "Credential mode:"
    echo "  1) Same username/password for all services"
    echo "  2) Different credentials per service"
    read -r -p "Select [1]: " cred_mode
    if [[ "${cred_mode:-1}" != "2" ]]; then
        shared_creds=true
        read -r -p "Username for all services [admin]: " shared_user
        shared_user=${shared_user:-admin}
        # Create a temporary shared htpasswd file — will be copied per service
        shared_htpasswd=$(mktemp)
        htpasswd -c "$shared_htpasswd" "$shared_user"
    fi

    # Create directories
    mkdir -p "$AUTH_DIR" "$SNIPPET_DIR" "$BACKUP_DIR"

    echo ""

    # Set up auth for each detected service
    for entry in "${SERVICES[@]}"; do
        local name config_file port
        name=$(get_service_name "$entry")
        config_file=$(get_config_file "$entry")
        port=$(get_service_port "$entry")
        local full_path="${CONFIG_DIR}/${config_file}"

        if [ ! -f "$full_path" ]; then
            continue
        fi

        local display_name
        display_name="${name^}"

        echo "--- ${display_name} ---"

        # Create htpasswd file
        if [ "$shared_creds" = true ]; then
            if [ -f "${AUTH_DIR}/.htpasswd-${name}" ]; then
                read -r -p "  htpasswd file exists. Overwrite with shared credentials? [Y/n]: " overwrite
                if [[ "$overwrite" =~ ^[Nn]$ ]]; then
                    echo "  Keeping existing htpasswd"
                else
                    cp "$shared_htpasswd" "${AUTH_DIR}/.htpasswd-${name}"
                    echo "  Using shared credentials"
                fi
            else
                cp "$shared_htpasswd" "${AUTH_DIR}/.htpasswd-${name}"
                echo "  Using shared credentials"
            fi
        else
            if [ -f "${AUTH_DIR}/.htpasswd-${name}" ]; then
                read -r -p "  htpasswd file exists. Recreate? [y/N]: " recreate
                if [[ ! "$recreate" =~ ^[Yy]$ ]]; then
                    echo "  Keeping existing htpasswd"
                else
                    read -r -p "  Username [admin]: " username
                    username=${username:-admin}
                    htpasswd -c "${AUTH_DIR}/.htpasswd-${name}" "$username"
                fi
            else
                read -r -p "  Username [admin]: " username
                username=${username:-admin}
                htpasswd -c "${AUTH_DIR}/.htpasswd-${name}" "$username"
            fi
        fi

        # Create auth snippet
        create_snippet "$name" "$display_name"

        # Insert into main location/proxy block
        insert_auth "$full_path" "$name"

        # Handle push endpoints for Mimir and Loki
        if [[ "$name" == "mimir" ]] || [[ "$name" == "loki" ]]; then
            if [ "$protect_push" = true ]; then
                insert_push_auth "$full_path" "$name"
            else
                echo "  ⚠ Push endpoint left open — consider IP restrictions"
            fi
        fi

        echo ""
    done

    # Clean up shared temp file
    [ -n "$shared_htpasswd" ] && rm -f "$shared_htpasswd"

    # Set permissions on htpasswd files
    chmod 640 "${AUTH_DIR}"/.htpasswd-* 2>/dev/null || true
    chown "root:${WEB_USER}" "${AUTH_DIR}"/.htpasswd-* 2>/dev/null || true

    # Test and reload
    echo "Testing ${WEB_SERVER} configuration..."
    if test_config; then
        systemctl reload "$SERVICE_NAME"
        echo ""
        echo "=========================================="
        echo "HTTP Basic Auth Successfully Configured!"
        echo "=========================================="
        echo ""
        echo "Web server: ${WEB_SERVER}"
        echo "Backups:    ${BACKUP_DIR}"
        echo ""
        echo "To remove auth later: $0 --remove"
        echo "To check status:      $0 --status"
    else
        echo ""
        echo "${WEB_SERVER} configuration test FAILED!"
        echo "Your backups are in ${BACKUP_DIR} — restore manually if needed."
        exit 1
    fi
}

# ============================================================================
# MAIN
# ============================================================================

main() {
    if [[ $EUID -ne 0 ]]; then
        die "This script must be run as root"
    fi

    case "${1:-}" in
        -h|--help) show_usage ;;
        --remove) do_remove ;;
        --status) show_status ;;
        *)
            install_htpasswd
            setup_auth
            ;;
    esac
}

main "$@"