#!/usr/bin/env bash

#########################################################################################
#### sbom-scanner.sh — Generate SBOMs and scan for vulnerabilities with Syft + Grype ####
#### Wrapper for container images, directories, and archives with CI-friendly output ####
#### Requires: bash 4+, optionally jq                                                ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./sbom-scanner.sh --scan nginx:latest                                         ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
RUN_MODE=""
TARGET=""
FORMAT="${SBOM_FORMAT:-table}"
OUTPUT_FILE="${SBOM_OUTPUT_FILE:-}"
MIN_SEVERITY="${SBOM_MIN_SEVERITY:-}"
FAIL_ON="${SBOM_FAIL_ON:-}"
IGNORE_FILE="${SBOM_IGNORE_FILE:-}"
SYFT_BIN="${SYFT_PATH:-syft}"
GRYPE_BIN="${GRYPE_PATH:-grype}"
INSTALL_DIR="${INSTALL_DIR:-/usr/local/bin}"
DB_UPDATE="${DB_UPDATE:-false}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
START_TIME=""
CRIT_COUNT=0
HIGH_COUNT=0
MED_COUNT=0
LOW_COUNT=0
NEG_COUNT=0
TOTAL_VULNS=0

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        # shellcheck disable=SC2034
        RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        # shellcheck disable=SC2034
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    else
        # shellcheck disable=SC2034
        RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${DIM}[DEBUG]${RESET} $*"; fi; }

die() {
    err "$*"
    exit 1
}

elapsed() {
    local end_time
    end_time=$(date +%s)
    echo "$(( end_time - START_TIME ))s"
}

# ── Dependency checks ────────────────────────────────────────────────
check_tool() {
    local tool="$1"
    local bin_var="$2"
    if ! command -v "$bin_var" &>/dev/null; then
        err "${tool} is not installed (looked for '${bin_var}')"
        err "Install with: ./$(basename "$0") --install"
        exit 1
    fi
    verbose "${tool} found at $(command -v "$bin_var")"
}

# ── Detect target type ───────────────────────────────────────────────
detect_target_type() {
    local t="$1"
    if [[ "$t" == dir:* ]]; then
        echo "directory"
    elif [[ -f "$t" ]]; then
        echo "archive"
    else
        echo "container image"
    fi
}

# ── Severity level to number ─────────────────────────────────────────
severity_to_num() {
    case "${1,,}" in
        critical)    echo 5 ;;
        high)        echo 4 ;;
        medium)      echo 3 ;;
        low)         echo 2 ;;
        negligible)  echo 1 ;;
        *)           echo 0 ;;
    esac
}

severity_color() {
    case "${1,,}" in
        critical)    echo "$RED" ;;
        high)        echo "$RED" ;;
        medium)      echo "$YELLOW" ;;
        low)         echo "$DIM" ;;
        negligible)  echo "$DIM" ;;
        *)           echo "$RESET" ;;
    esac
}

# ══════════════════════════════════════════════════════════════════════
#                          INSTALL MODE
# ══════════════════════════════════════════════════════════════════════

do_install() {
    log "Installing Syft and Grype to ${INSTALL_DIR}..."

    if ! command -v curl &>/dev/null; then
        die "curl is required for installation"
    fi

    # Install Syft
    if command -v "$SYFT_BIN" &>/dev/null; then
        log "Syft already installed: $(command -v "$SYFT_BIN")"
    else
        log "Installing Syft..."
        verbose "curl -sSfL https://raw.githubusercontent.com/anchore/syft/main/install.sh | sh -s -- -b ${INSTALL_DIR}"
        if curl -sSfL https://raw.githubusercontent.com/anchore/syft/main/install.sh | sh -s -- -b "$INSTALL_DIR" 2>/dev/null; then
            echo -e "  ${GREEN}✓${RESET} Syft installed to ${INSTALL_DIR}/syft"
        else
            err "Failed to install Syft"
        fi
    fi

    # Install Grype
    if command -v "$GRYPE_BIN" &>/dev/null; then
        log "Grype already installed: $(command -v "$GRYPE_BIN")"
    else
        log "Installing Grype..."
        verbose "curl -sSfL https://raw.githubusercontent.com/anchore/grype/main/install.sh | sh -s -- -b ${INSTALL_DIR}"
        if curl -sSfL https://raw.githubusercontent.com/anchore/grype/main/install.sh | sh -s -- -b "$INSTALL_DIR" 2>/dev/null; then
            echo -e "  ${GREEN}✓${RESET} Grype installed to ${INSTALL_DIR}/grype"
        else
            err "Failed to install Grype"
        fi
    fi

    log "Installation complete"
}

# ══════════════════════════════════════════════════════════════════════
#                          SBOM MODE
# ══════════════════════════════════════════════════════════════════════

do_sbom() {
    check_tool "Syft" "$SYFT_BIN"

    local target_type
    target_type=$(detect_target_type "$TARGET")
    log "Generating SBOM with Syft..."
    log "Target: ${TARGET} (${target_type})"

    local syft_args=("$TARGET")

    case "$FORMAT" in
        table)           syft_args+=(-o syft-table) ;;
        json)            syft_args+=(-o syft-json) ;;
        cyclonedx-json)  syft_args+=(-o cyclonedx-json) ;;
        *)               syft_args+=(-o syft-table) ;;
    esac

    local output
    if output=$("$SYFT_BIN" "${syft_args[@]}" 2>/dev/null); then
        if [[ -n "$OUTPUT_FILE" ]]; then
            echo "$output" > "$OUTPUT_FILE"
            log "Output written to ${OUTPUT_FILE}"
        else
            echo ""
            echo "$output"
        fi
    else
        die "Syft failed to generate SBOM for ${TARGET}"
    fi

    # Count packages for summary
    if [[ "$FORMAT" == "table" ]]; then
        local pkg_count
        pkg_count=$(echo "$output" | tail -n +2 | grep -c "" || true)
        echo ""
        log "Total packages: ${pkg_count}"
    fi

    log "Generated in $(elapsed)"
}

# ══════════════════════════════════════════════════════════════════════
#                          SCAN MODE
# ══════════════════════════════════════════════════════════════════════

do_scan() {
    check_tool "Grype" "$GRYPE_BIN"

    local target_type
    target_type=$(detect_target_type "$TARGET")
    log "Scanning for vulnerabilities with Grype..."
    log "Target: ${TARGET} (${target_type})"

    if [[ "$DB_UPDATE" == "true" ]]; then
        log "Updating vulnerability database..."
        "$GRYPE_BIN" db update 2>/dev/null || warn "Database update failed, using cached DB"
    fi

    local grype_args=("$TARGET" -o json)
    [[ -n "$IGNORE_FILE" ]] && grype_args+=(--config "$IGNORE_FILE")

    local raw_json
    if ! raw_json=$("$GRYPE_BIN" "${grype_args[@]}" 2>/dev/null); then
        die "Grype failed to scan ${TARGET}"
    fi

    # Parse vulnerability counts
    if command -v jq &>/dev/null; then
        CRIT_COUNT=$(echo "$raw_json" | jq '[.matches[]? | select(.vulnerability.severity == "Critical")] | length')
        HIGH_COUNT=$(echo "$raw_json" | jq '[.matches[]? | select(.vulnerability.severity == "High")] | length')
        MED_COUNT=$(echo "$raw_json" | jq '[.matches[]? | select(.vulnerability.severity == "Medium")] | length')
        LOW_COUNT=$(echo "$raw_json" | jq '[.matches[]? | select(.vulnerability.severity == "Low")] | length')
        NEG_COUNT=$(echo "$raw_json" | jq '[.matches[]? | select(.vulnerability.severity == "Negligible")] | length')
        TOTAL_VULNS=$(echo "$raw_json" | jq '[.matches[]?] | length')
    else
        TOTAL_VULNS=$(echo "$raw_json" | grep -c '"vulnerability"' || true)
    fi

    local min_sev_num=0
    [[ -n "$MIN_SEVERITY" ]] && min_sev_num=$(severity_to_num "$MIN_SEVERITY")

    case "$FORMAT" in
        json)
            if [[ -n "$OUTPUT_FILE" ]]; then
                echo "$raw_json" > "$OUTPUT_FILE"
                log "JSON output written to ${OUTPUT_FILE}"
            else
                echo "$raw_json"
            fi
            ;;
        prometheus)
            print_prometheus_metrics
            ;;
        *)
            print_scan_table "$raw_json" "$min_sev_num"
            ;;
    esac

    print_scan_summary

    # Check fail threshold
    check_fail_threshold
}

print_scan_table() {
    local raw_json="$1"
    local min_sev_num="$2"

    if ! command -v jq &>/dev/null; then
        warn "jq not installed — cannot format table output, showing raw"
        echo "$raw_json"
        return
    fi

    echo ""
    printf "  ${BOLD}%-16s %-10s %-20s %-14s %s${RESET}\n" "VULNERABILITY" "SEVERITY" "PACKAGE" "VERSION" "FIXED-IN"
    printf "  %s\n" "$(printf '%.0s─' {1..75})"

    echo "$raw_json" | jq -c '.matches[]?' | while IFS= read -r match; do
        local vuln_id severity pkg_name pkg_ver fixed_in
        vuln_id=$(echo "$match" | jq -r '.vulnerability.id')
        severity=$(echo "$match" | jq -r '.vulnerability.severity')
        pkg_name=$(echo "$match" | jq -r '.artifact.name')
        pkg_ver=$(echo "$match" | jq -r '.artifact.version')
        fixed_in=$(echo "$match" | jq -r '.vulnerability.fix.versions[0] // "(none)"')

        local sev_num
        sev_num=$(severity_to_num "$severity")
        [[ "$sev_num" -lt "$min_sev_num" ]] && continue

        local color
        color=$(severity_color "$severity")
        printf "  %-16s ${color}%-10s${RESET} %-20s %-14s %s\n" \
            "$vuln_id" "$severity" "${pkg_name:0:20}" "${pkg_ver:0:14}" "$fixed_in"
    done

    if [[ -n "$OUTPUT_FILE" ]]; then
        echo "$raw_json" > "$OUTPUT_FILE"
        log "Full JSON written to ${OUTPUT_FILE}"
    fi
}

print_scan_summary() {
    echo ""
    echo -e "  ${BOLD}Summary${RESET}"
    echo -e "    Total vulnerabilities: ${TOTAL_VULNS}"
    [[ "$CRIT_COUNT" -gt 0 ]] && echo -e "      ${RED}Critical: ${CRIT_COUNT}${RESET}" || echo "      Critical: 0"
    [[ "$HIGH_COUNT" -gt 0 ]] && echo -e "      ${RED}High:     ${HIGH_COUNT}${RESET}" || echo "      High:     0"
    [[ "$MED_COUNT" -gt 0 ]]  && echo -e "      ${YELLOW}Medium:   ${MED_COUNT}${RESET}" || echo "      Medium:   0"
    echo "      Low:      ${LOW_COUNT}"
    echo "      Negligible: ${NEG_COUNT}"
    log "Scanned in $(elapsed)"
}

print_prometheus_metrics() {
    local ts
    ts=$(date +%s)
    cat <<EOF
# HELP sbom_scan_vulnerabilities_total Total vulnerabilities found
# TYPE sbom_scan_vulnerabilities_total gauge
sbom_scan_vulnerabilities_total{target="${TARGET}"} ${TOTAL_VULNS}
# HELP sbom_scan_vulnerabilities_by_severity Vulnerabilities by severity
# TYPE sbom_scan_vulnerabilities_by_severity gauge
sbom_scan_vulnerabilities_by_severity{target="${TARGET}",severity="critical"} ${CRIT_COUNT}
sbom_scan_vulnerabilities_by_severity{target="${TARGET}",severity="high"} ${HIGH_COUNT}
sbom_scan_vulnerabilities_by_severity{target="${TARGET}",severity="medium"} ${MED_COUNT}
sbom_scan_vulnerabilities_by_severity{target="${TARGET}",severity="low"} ${LOW_COUNT}
sbom_scan_vulnerabilities_by_severity{target="${TARGET}",severity="negligible"} ${NEG_COUNT}
# HELP sbom_scan_timestamp_seconds Scan timestamp
# TYPE sbom_scan_timestamp_seconds gauge
sbom_scan_timestamp_seconds{target="${TARGET}"} ${ts}
EOF
}

check_fail_threshold() {
    if [[ -z "$FAIL_ON" ]]; then
        return
    fi

    local threshold_num
    threshold_num=$(severity_to_num "$FAIL_ON")
    local exceeded=0

    if [[ "$threshold_num" -le 5 && "$CRIT_COUNT" -gt 0 ]]; then exceeded=1; fi
    if [[ "$threshold_num" -le 4 && "$HIGH_COUNT" -gt 0 ]]; then exceeded=1; fi
    if [[ "$threshold_num" -le 3 && "$MED_COUNT" -gt 0 ]]; then exceeded=1; fi
    if [[ "$threshold_num" -le 2 && "$LOW_COUNT" -gt 0 ]]; then exceeded=1; fi
    if [[ "$threshold_num" -le 1 && "$NEG_COUNT" -gt 0 ]]; then exceeded=1; fi

    if [[ "$exceeded" -eq 1 ]]; then
        echo ""
        err "Vulnerabilities at or above '${FAIL_ON}' severity found — failing"
        exit 2
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                          FULL MODE
# ══════════════════════════════════════════════════════════════════════

do_full() {
    check_tool "Syft" "$SYFT_BIN"
    check_tool "Grype" "$GRYPE_BIN"

    local target_type
    target_type=$(detect_target_type "$TARGET")
    log "Running full SBOM + vulnerability scan..."
    log "Target: ${TARGET} (${target_type})"

    if [[ "$DB_UPDATE" == "true" ]]; then
        log "Updating vulnerability database..."
        "$GRYPE_BIN" db update 2>/dev/null || warn "Database update failed, using cached DB"
    fi

    # Generate SBOM
    log "Phase 1: Generating SBOM with Syft..."
    local sbom_json
    if ! sbom_json=$("$SYFT_BIN" "$TARGET" -o syft-json 2>/dev/null); then
        die "Syft failed to generate SBOM"
    fi

    local pkg_count
    if command -v jq &>/dev/null; then
        pkg_count=$(echo "$sbom_json" | jq '[.artifacts[]?] | length')
    else
        pkg_count="?"
    fi
    log "SBOM complete: ${pkg_count} packages found"

    # Scan with Grype using the SBOM as input
    log "Phase 2: Scanning for vulnerabilities with Grype..."
    local grype_args=(--sbom -)
    grype_args+=(-o json)
    [[ -n "$IGNORE_FILE" ]] && grype_args+=(--config "$IGNORE_FILE")

    local raw_json
    if ! raw_json=$(echo "$sbom_json" | "$GRYPE_BIN" "${grype_args[@]}" 2>/dev/null); then
        die "Grype scan failed"
    fi

    # Parse counts
    if command -v jq &>/dev/null; then
        CRIT_COUNT=$(echo "$raw_json" | jq '[.matches[]? | select(.vulnerability.severity == "Critical")] | length')
        HIGH_COUNT=$(echo "$raw_json" | jq '[.matches[]? | select(.vulnerability.severity == "High")] | length')
        MED_COUNT=$(echo "$raw_json" | jq '[.matches[]? | select(.vulnerability.severity == "Medium")] | length')
        LOW_COUNT=$(echo "$raw_json" | jq '[.matches[]? | select(.vulnerability.severity == "Low")] | length')
        NEG_COUNT=$(echo "$raw_json" | jq '[.matches[]? | select(.vulnerability.severity == "Negligible")] | length')
        TOTAL_VULNS=$(echo "$raw_json" | jq '[.matches[]?] | length')
    fi

    local min_sev_num=0
    [[ -n "$MIN_SEVERITY" ]] && min_sev_num=$(severity_to_num "$MIN_SEVERITY")

    case "$FORMAT" in
        json)
            if [[ -n "$OUTPUT_FILE" ]]; then
                echo "$raw_json" > "$OUTPUT_FILE"
                log "JSON output written to ${OUTPUT_FILE}"
            else
                echo "$raw_json"
            fi
            ;;
        prometheus)
            print_prometheus_metrics
            ;;
        *)
            print_scan_table "$raw_json" "$min_sev_num"
            ;;
    esac

    print_scan_summary
    check_fail_threshold
}

# ══════════════════════════════════════════════════════════════════════
#                          HELP
# ══════════════════════════════════════════════════════════════════════

show_help() {
    cat <<EOF
Usage: $SCRIPT_NAME [MODE] TARGET [OPTIONS]

Generate SBOMs and scan for vulnerabilities using Syft and Grype.

MODES:
    --sbom TARGET           Generate SBOM for target
    --scan TARGET           Scan target for vulnerabilities
    --full TARGET           Generate SBOM and scan in one pass
    --install               Install Syft and Grype

OPTIONS:
    --format FORMAT         Output: table (default), json, cyclonedx-json, prometheus
    --output-file FILE      Write output to file instead of stdout
    --min-severity SEV      Minimum severity to display (negligible, low, medium, high, critical)
    --fail-on SEV           Exit with code 2 if vulnerabilities at or above severity found
    --ignore-file FILE      Path to Grype ignore file for suppressing false positives
    --db-update             Force update vulnerability databases before scanning
    --verbose               Debug output
    --no-color              Disable colored output
    --help, -h              Show this help

TARGETS:
    nginx:latest            Container image (from registry or local)
    dir:/path/to/project    Local directory
    /path/to/image.tar      Archive file

ENVIRONMENT VARIABLES:
    SBOM_FORMAT             Output format (default: table)
    SBOM_OUTPUT_FILE        Write output to this file
    SBOM_MIN_SEVERITY       Minimum severity to display
    SBOM_FAIL_ON            Fail threshold severity
    SBOM_IGNORE_FILE        Path to Grype ignore rules
    SYFT_PATH               Path to Syft binary (default: syft)
    GRYPE_PATH              Path to Grype binary (default: grype)
    INSTALL_DIR             Install directory (default: /usr/local/bin)
    DB_UPDATE               Force database update (default: false)
    VERBOSE                 Debug output (default: false)
    NO_COLOR                Disable colored output (default: false)

EXIT CODES:
    0   Success
    1   Runtime error
    2   Vulnerability threshold exceeded (--fail-on)

EXAMPLES:
    # Generate SBOM for a container image
    ./$SCRIPT_NAME --sbom nginx:latest

    # Scan for vulnerabilities, show high and critical only
    ./$SCRIPT_NAME --scan nginx:latest --min-severity high

    # CI mode — fail if critical vulns found
    ./$SCRIPT_NAME --scan nginx:latest --fail-on critical

    # Full SBOM + scan with JSON output
    ./$SCRIPT_NAME --full nginx:latest --format json --output-file report.json

    # Install Syft and Grype
    ./$SCRIPT_NAME --install
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                          MAIN
# ══════════════════════════════════════════════════════════════════════

main() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --sbom)
                RUN_MODE="sbom"
                shift
                if [[ $# -gt 0 && ! "$1" =~ ^-- ]]; then
                    TARGET="$1"; shift
                fi
                ;;
            --scan)
                RUN_MODE="scan"
                shift
                if [[ $# -gt 0 && ! "$1" =~ ^-- ]]; then
                    TARGET="$1"; shift
                fi
                ;;
            --full)
                RUN_MODE="full"
                shift
                if [[ $# -gt 0 && ! "$1" =~ ^-- ]]; then
                    TARGET="$1"; shift
                fi
                ;;
            --install)
                RUN_MODE="install"
                shift
                ;;
            --format)
                FORMAT="$2"; shift 2 ;;
            --output-file)
                OUTPUT_FILE="$2"; shift 2 ;;
            --min-severity)
                MIN_SEVERITY="$2"; shift 2 ;;
            --fail-on)
                FAIL_ON="$2"; shift 2 ;;
            --ignore-file)
                IGNORE_FILE="$2"; shift 2 ;;
            --db-update)
                DB_UPDATE="true"; shift ;;
            --verbose)
                VERBOSE="true"; shift ;;
            --no-color)
                COLOR="never"; shift ;;
            --help|-h)
                show_help; exit 0 ;;
            *)
                if [[ -z "$TARGET" && ! "$1" =~ ^-- ]]; then
                    TARGET="$1"; shift
                else
                    die "Unknown option: $1 (see --help)"
                fi
                ;;
        esac
    done

    # Handle NO_COLOR env var
    if [[ "${NO_COLOR:-}" == "true" ]]; then
        COLOR="never"
    fi

    setup_colors

    if [[ -z "$RUN_MODE" ]]; then err "No mode specified"; echo ""; show_help; exit 1; fi

    if [[ "$RUN_MODE" != "install" && -z "$TARGET" ]]; then
        die "No target specified. Provide a container image, dir:path, or archive path"
    fi

    START_TIME=$(date +%s)

    echo ""
    echo -e "${BOLD}SBOM Scanner${RESET}"
    if [[ -n "$TARGET" ]]; then
        echo "Target: ${TARGET}"
    fi
    echo "Mode:   ${RUN_MODE}"
    echo "Format: ${FORMAT}"
    echo "Time:   $(date -u '+%Y-%m-%dT%H:%M:%SZ')"
    echo ""

    case "$RUN_MODE" in
        sbom)    do_sbom ;;
        scan)    do_scan ;;
        full)    do_full ;;
        install) do_install ;;
    esac
}

main "$@"