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

# Container Image Security Scanner
# Uses Trivy to scan images for CVEs

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VAULT_PASSWORD_FILE="${VAULT_PASSWORD_FILE:-secrets/.vault_pass}"
REPORT_DIR="${REPORT_DIR:-/var/log/trivy-scans}"
SEVERITY="${SEVERITY:-HIGH,CRITICAL}"

usage() {
    echo "Usage: $0 [options]"
    echo "Options:"
    echo "  --scan-all       Scan all running containers (default)"
    echo "  --scan-image IMAGE   Scan specific image"
    echo "  --severity LEVEL     Severity filter (default: HIGH,CRITICAL)"
    echo "  --report-only        Exit 0 even if CVEs found (for CI reporting)"
    echo "  --help               Show this help"
    exit 1
}

SCAN_MODE="all"
TARGET_IMAGE=""
REPORT_ONLY=false

while [[ $# -gt 0 ]]; do
    case $1 in
        --scan-all) SCAN_MODE="all"; shift ;;
        --scan-image)
            SCAN_MODE="single"
            TARGET_IMAGE="$2"
            shift 2
            ;;
        --severity) SEVERITY="$2"; shift 2 ;;
        --report-only) REPORT_ONLY=true; shift ;;
        --help) usage ;;
        *) echo "Unknown option: $1"; usage ;;
    esac
done

# Check if Trivy is installed
if ! command -v trivy &> /dev/null; then
    echo "Installing Trivy..."
    curl -sfL https://raw.githubusercontent.com/aquasecurity/trivy/main/contrib/install.sh | sh -s -- -b /usr/local/bin
fi

# Create report directory
mkdir -p "$REPORT_DIR"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
REPORT_FILE="$REPORT_DIR/trivy-scan-$TIMESTAMP.json"
SUMMARY_FILE="$REPORT_DIR/trivy-scan-$TIMESTAMP.summary"

scan_image() {
    local image="$1"
    echo "Scanning: $image"
    
    # Scan with Trivy
    if trivy image \
        --severity "$SEVERITY" \
        --format json \
        --output "$REPORT_FILE.tmp" \
        --scanners vuln,config \
        "$image" 2>/dev/null; then
        
        # Check for vulnerabilities
        if command -v jq &> /dev/null; then
            CRITICAL=$(jq '[.Results[]?.Vulnerabilities? // empty | .[] | select(.Severity == "CRITICAL")] | length' "$REPORT_FILE.tmp" 2>/dev/null || echo "0")
            HIGH=$(jq '[.Results[]?.Vulnerabilities? // empty | .[] | select(.Severity == "HIGH")] | length' "$REPORT_FILE.tmp" 2>/dev/null || echo "0")
            
            echo "  CRITICAL: $CRITICAL, HIGH: $HIGH"
            
            if [[ "$CRITICAL" -gt 0 ]] && [[ "$REPORT_ONLY" == "false" ]]; then
                echo "ERROR: CRITICAL vulnerabilities found in $image" >&2
                return 1
            fi
        fi
        
        cat "$REPORT_FILE.tmp" >> "$REPORT_FILE"
        rm -f "$REPORT_FILE.tmp"
        return 0
    else
        echo "WARNING: Failed to scan $image" >&2
        return 0  # Continue scanning other images
    fi
}

# Main scanning logic
echo "Starting Trivy container scan at $(date)"
echo "Severity filter: $SEVERITY"
echo ""

if [[ "$SCAN_MODE" == "single" ]]; then
    scan_image "$TARGET_IMAGE"
else
    # Get all unique images from running containers
    IMAGES=$(docker ps --format '{{.Image}}' | sort -u)
    
    TOTAL=0
    FAILED=0
    
    for image in $IMAGES; do
        TOTAL=$((TOTAL + 1))
        if ! scan_image "$image"; then
            FAILED=$((FAILED + 1))
        fi
    done
    
    echo ""
    echo "Scan complete: $TOTAL images scanned, $FAILED failed"
fi

# Generate summary
echo ""
echo "Report saved to: $REPORT_FILE"
echo "To view: trivy convert --format table $REPORT_FILE"

# Cleanup old reports (keep last 30 days)
find "$REPORT_DIR" -name "trivy-scan-*.json" -mtime +30 -delete 2>/dev/null || true

if [[ "$FAILED" -gt 0 ]] && [[ "$REPORT_ONLY" == "false" ]]; then
    exit 1
fi

exit 0