#!/bin/bash
################################################################################
# Script Name: trivy-cve-auditor.sh
# Version: 1.0
# Description: Prometheus exporter that scans all local container images with
#              Trivy and outputs vulnerability metrics by severity. Supports
#              stdout, node_exporter textfile collector, and HTTP server modes.
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Prerequisites:
#   - trivy installed and in PATH
#   - docker or podman (auto-detected)
#   - jq for JSON parsing
#   - nc (netcat) for HTTP mode
#
# Usage:
#   # Output to stdout (default)
#   ./trivy-cve-auditor.sh
#
#   # Write to node_exporter textfile collector
#   ./trivy-cve-auditor.sh --textfile
#
#   # HTTP server mode
#   ./trivy-cve-auditor.sh --http
#
# Environment Variables:
#   TRIVY_SEVERITY   Severity levels to scan (default: HIGH,CRITICAL)
#   SKIP_IMAGES      Regex pattern to exclude images
#   TRIVY_TIMEOUT    Per-image scan timeout in seconds (default: 300)
#   PROM_PORT        HTTP listen port for --http mode (default: 9199)
#
################################################################################

set -euo pipefail

# ============================================================================
# CONFIGURATION
# ============================================================================

TRIVY_SEVERITY="${TRIVY_SEVERITY:-HIGH,CRITICAL}"
SKIP_IMAGES="${SKIP_IMAGES:-}"
TRIVY_TIMEOUT="${TRIVY_TIMEOUT:-300}"
PROM_PORT="${PROM_PORT:-9199}"
TEXTFILE_DIR="/var/lib/node_exporter/textfile"
TEXTFILE_PATH="${TEXTFILE_DIR}/trivy_cve.prom"

# Runtime
MODE="stdout"
CONTAINER_CMD=""
FIRST_SCAN=true

# ============================================================================
# HELPERS
# ============================================================================

log() { echo "# $*" >&2; }

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

Scan all local container images with Trivy and export Prometheus metrics.

MODES:
    (default)       Print metrics to stdout
    --textfile      Write to node_exporter textfile collector
    --http          Serve metrics on HTTP port (default: $PROM_PORT)

OPTIONS:
    -h, --help      Show this help message

ENVIRONMENT VARIABLES:
    TRIVY_SEVERITY   Severity levels (default: HIGH,CRITICAL)
    SKIP_IMAGES      Regex to exclude images from scanning
    TRIVY_TIMEOUT    Per-image timeout in seconds (default: 300)
    PROM_PORT        HTTP port for --http mode (default: 9199)

EXAMPLES:
    $(basename "$0")                                    # stdout
    $(basename "$0") --textfile                         # textfile collector
    $(basename "$0") --http                             # HTTP server on :9199
    PROM_PORT=9200 $(basename "$0") --http              # custom port
    SKIP_IMAGES="pause|kindest" $(basename "$0")        # skip images

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --textfile) MODE="textfile"; shift ;;
            --http)     MODE="http"; shift ;;
            -h|--help)  show_usage ;;
            *) log "Unknown option: $1"; exit 1 ;;
        esac
    done
}

detect_container_runtime() {
    if command -v podman &>/dev/null; then
        CONTAINER_CMD="podman"
    elif command -v docker &>/dev/null; then
        CONTAINER_CMD="docker"
    else
        log "ERROR: neither docker nor podman found in PATH"
        exit 1
    fi
    log "Detected container runtime: $CONTAINER_CMD"
}

check_dependencies() {
    if ! command -v trivy &>/dev/null; then
        log "ERROR: trivy not found in PATH"
        exit 1
    fi
    if ! command -v jq &>/dev/null; then
        log "ERROR: jq not found in PATH"
        exit 1
    fi
    if [[ "$MODE" == "http" ]] && ! command -v nc &>/dev/null; then
        log "ERROR: nc (netcat) required for HTTP mode"
        exit 1
    fi
}

# Get list of local images as "repository:tag" one per line
get_local_images() {
    $CONTAINER_CMD images --format '{{.Repository}}:{{.Tag}}' 2>/dev/null | \
        grep -v '<none>' | \
        sort -u
}

# ============================================================================
# METRICS GENERATION
# ============================================================================

generate_metrics() {
    local scan_start
    scan_start=$(date +%s)

    local images_scanned=0
    local images_with_critical=0

    local image_list
    image_list=$(get_local_images)

    if [[ -z "$image_list" ]]; then
        log "No local images found"
    fi

    # HELP/TYPE headers
    cat <<'HEADER'
# HELP trivy_image_vulnerabilities Number of vulnerabilities per image per severity
# TYPE trivy_image_vulnerabilities gauge
# HELP trivy_image_last_scan_timestamp Unix timestamp of last scan per image
# TYPE trivy_image_last_scan_timestamp gauge
HEADER

    local db_flag=""

    while IFS= read -r image; do
        [[ -z "$image" ]] && continue

        # Skip images matching exclusion pattern
        if [[ -n "$SKIP_IMAGES" ]] && echo "$image" | grep -qE "$SKIP_IMAGES"; then
            log "Skipping excluded image: $image"
            continue
        fi

        log "Scanning: $image"

        # After first scan, skip DB update to save time
        if [[ "$FIRST_SCAN" == true ]]; then
            db_flag=""
            FIRST_SCAN=false
        else
            db_flag="--skip-db-update"
        fi

        local json_output
        if ! json_output=$(trivy image \
            --format json \
            --severity "$TRIVY_SEVERITY" \
            --timeout "${TRIVY_TIMEOUT}s" \
            --quiet \
            $db_flag \
            "$image" 2>/dev/null); then
            log "WARN: failed to scan $image"
            continue
        fi

        # Parse vulnerability counts by severity
        local high_count=0
        local critical_count=0

        if echo "$json_output" | jq -e '.Results' &>/dev/null; then
            high_count=$(echo "$json_output" | \
                jq '[.Results[]? | .Vulnerabilities[]? | select(.Severity == "HIGH")] | length')
            critical_count=$(echo "$json_output" | \
                jq '[.Results[]? | .Vulnerabilities[]? | select(.Severity == "CRITICAL")] | length')
        fi

        # Sanitize image name for Prometheus label (replace special chars)
        local label_image
        label_image="${image//\"/\\\"}"

        local now
        now=$(date +%s)

        echo "trivy_image_vulnerabilities{image=\"${label_image}\",severity=\"HIGH\"} ${high_count}"
        echo "trivy_image_vulnerabilities{image=\"${label_image}\",severity=\"CRITICAL\"} ${critical_count}"
        echo "trivy_image_last_scan_timestamp{image=\"${label_image}\"} ${now}"

        images_scanned=$((images_scanned + 1))
        if [[ "$critical_count" -gt 0 ]]; then
            images_with_critical=$((images_with_critical + 1))
        fi

    done <<< "$image_list"

    echo ""

    # Summary metrics
    local scan_end
    scan_end=$(date +%s)
    local duration=$((scan_end - scan_start))

    cat <<EOF

# HELP trivy_scan_duration_seconds Total time to scan all images
# TYPE trivy_scan_duration_seconds gauge
trivy_scan_duration_seconds ${duration}

# HELP trivy_images_scanned_total Number of images scanned
# TYPE trivy_images_scanned_total gauge
trivy_images_scanned_total ${images_scanned}

# HELP trivy_images_with_critical_total Images with at least one CRITICAL CVE
# TYPE trivy_images_with_critical_total gauge
trivy_images_with_critical_total ${images_with_critical}
EOF
}

# ============================================================================
# OUTPUT MODES
# ============================================================================

run_stdout() {
    generate_metrics
}

run_textfile() {
    if [[ ! -d "$TEXTFILE_DIR" ]]; then
        log "Creating textfile directory: $TEXTFILE_DIR"
        mkdir -p "$TEXTFILE_DIR"
    fi

    local temp_file
    temp_file=$(mktemp "${TEXTFILE_DIR}/.trivy_cve.XXXXXX")

    if ! generate_metrics > "$temp_file" 2>/dev/null; then
        rm -f "$temp_file"
        log "ERROR: failed to generate metrics"
        exit 1
    fi

    chmod 644 "$temp_file"
    mv -f "$temp_file" "$TEXTFILE_PATH"
    log "Metrics written to $TEXTFILE_PATH"
}

run_http() {
    log "Starting Trivy CVE Auditor on port $PROM_PORT..."

    if ! command -v nc &>/dev/null; then
        log "ERROR: nc (netcat) required for HTTP mode"
        exit 1
    fi

    while true; do
        {
            read -r request
            if [[ "$request" =~ ^GET\ /metrics ]]; then
                local body
                body=$(generate_metrics 2>/dev/null)
                local length=${#body}
                echo -e "HTTP/1.1 200 OK\r\nContent-Type: text/plain; version=0.0.4\r\nContent-Length: ${length}\r\n\r"
                echo "$body"
            else
                local html="<html><head><title>Trivy CVE Auditor</title></head><body><h1>Trivy CVE Auditor</h1><p><a href=\"/metrics\">Metrics</a></p></body></html>"
                local length=${#html}
                echo -e "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: ${length}\r\n\r"
                echo "$html"
            fi
        } | nc -l -p "$PROM_PORT" -q 1 2>/dev/null
    done
}

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

main() {
    parse_args "$@"
    check_dependencies
    detect_container_runtime

    case "$MODE" in
        stdout)   run_stdout ;;
        textfile) run_textfile ;;
        http)     run_http ;;
    esac
}

main "$@"