#!/usr/bin/env bash

# Copyright 2025 The CHPs-dev Authors
# SPDX-License-Identifier: Apache-2.0

# CHPs Scorer - Container Hardening Priorities Scoring Tool
# This script evaluates OCI container images against the CHPs criteria

# Displays help
display_help() {
    echo "CHPs Scorer"
    echo
    echo "Container Hardening Priorities Scoring Tool evaluates OCI container images"
    echo "against the Container Hardening Priorities (CHPs) criteria."
    echo
    echo "CHPs home page:       https://github.com/chps-dev/chps"
    echo "Project home page:    https://github.com/chps-dev/chps-scorer"
    echo
    echo "Usage:"
    echo "  $0 [options] <image>"
    echo
    echo "Options:"
    echo "  --skip-cves             Skip CVE checks, defaults to false"
    echo "  --dockerfile <path>     Use Dockerfile for additional checks"
    echo "  -o, --output <format>   Set output format: text (default), json, html, badges"
    echo "  --local                 Use local image instead of pulling from registry"
    echo "  -h, --help              Display this help message and exit"
}

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BOLD='\033[1m'
NC='\033[0m' # No Color

# Grade colors
A_PLUS_COLOR='\033[1;32m'  # Bright green
A_COLOR='\033[0;32m'       # Green
B_COLOR='\033[1;33m'       # Bright yellow
C_COLOR='\033[0;33m'       # Yellow
D_COLOR='\033[0;31m'       # Red
E_COLOR='\033[1;31m'       # Bright red

# Function to calculate grade based on score
get_grade() {

    # Played with these grades as many vectors only have 4 max points
    local score=$1
    local max_score=$2
    local percentage=$((score * 100 / max_score))

    if [ "$score" -eq 0 ]; then
        echo "E"
    elif [ "$score" -eq "$max_score" ]; then
        echo "A+"
    elif [ $percentage -ge 75 ]; then
        echo "A"
    elif [ $percentage -ge 50 ]; then
        echo "B"
    elif [ $percentage -ge 40 ]; then
        echo "C"
    else
        echo "D"
    fi
}

# Source the check modules
source "$(dirname "$0")/minimalism_checks.sh"
source "$(dirname "$0")/provenance_checks.sh"
source "$(dirname "$0")/config_checks.sh"
source "$(dirname "$0")/cve_checks.sh"

# Function to check if a command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Parse command line arguments
SKIP_CVES=false
DOCKERFILE=""
OUTPUT_FORMAT="text"
USE_LOCAL_IMAGE=false
while [[ $# -gt 0 ]]; do
    case $1 in
        --skip-cves)
            SKIP_CVES=true
            shift
            ;;
        -d|--dockerfile)
            DOCKERFILE="$2"
            shift 2
            ;;
        -o|--output)
            OUTPUT_FORMAT="$2"
            shift 2
            ;;
        -h|--help)
            display_help
            exit 0
            ;;
        --local)
            USE_LOCAL_IMAGE=true
            shift
            ;;
        *)
            break
            ;;
    esac
done

# Check if image name is provided
if [ $# -eq 0 ]; then
    display_help
    exit 1
fi

# Function to generate badge URL
get_badge_url() {
    local label=$1
    local grade=$2
    local label_color="%233443F4"  # Blue label background
    local color

    # Set color based on grade
    case $grade in
        "A+") color="%2301A178";; # Green
        "A")  color="%2304B45F";; # Light green
        "B")  color="%23FFB000";; # Yellow
        "C")  color="%23FF8C00";; # Orange
        "D")  color="%23FF4400";; # Light red
        "E")  color="%23FF0000";; # Red
        *)    color="%23808080";; # Gray for unknown
    esac

    # URL encode the grade (replace + with %2B)
    local encoded_grade=${grade//+/%2B}

    echo "https://img.shields.io/badge/${label}-${encoded_grade}-gold?style=flat-square&labelColor=${label_color}&color=${color}"
}

# Function to check if curl exists
check_curl() {
    if ! command -v curl &> /dev/null; then
        echo "curl is required for image display but not found" >&2
        return 1
    fi
    return 0
}

# Function to detect terminal image support
detect_term_img_support() {
    # Check for iTerm2
    if [[ -n "$ITERM_SESSION_ID" ]]; then
        if [[ -n "$TERM_PROGRAM_VERSION" ]]; then
            echo "iterm"
            return 0
        fi
    fi

    # Check for Kitty
    if [[ -n "$KITTY_WINDOW_ID" ]]; then
        echo "kitty"
        return 0
    fi

    # Check for terminals that support sixel
    if [[ "$TERM" =~ "xterm" ]] && command -v img2sixel &> /dev/null; then
        echo "sixel"
        return 0
    fi

    echo "none"
    return 1
}

# Function to display image in terminal
display_badge() {
    local url="$1"
    local term_type="$2"
    local tmp_file="/tmp/chps-badge-$RANDOM.png"

    # Download the badge
    if ! curl -s "$url" -o "$tmp_file"; then
        echo "Failed to download badge" >&2
        return 1
    fi

    case "$term_type" in
        "iterm")
            # For iTerm2, we can specify width in pixels or cells
            printf '\033]1337;File=inline=1;width=auto;height=auto:'
            base64 < "$tmp_file"
            printf '\a\n'
            ;;
        "kitty")
            # For Kitty, we can specify scale factor
            kitty +kitten icat --scale-up --transfer-mode=file --align=left "$tmp_file"
            ;;
        "sixel")
            # For Sixel, we can specify width
            img2sixel "$tmp_file"
            ;;
    esac

    rm -f "$tmp_file"
}

# Function to get grade color
get_grade_color() {
    local grade=$1
    case $grade in
        "A+") echo -e "${A_PLUS_COLOR}";;
        "A")  echo -e "${A_COLOR}";;
        "B")  echo -e "${B_COLOR}";;
        "C")  echo -e "${C_COLOR}";;
        "D")  echo -e "${D_COLOR}";;
        "E")  echo -e "${E_COLOR}";;
        *)    echo -e "${NC}";;
    esac
}
output_html() {
    local image=$1
    local digest=$2
    local minimalism_score=$3
    local provenance_score=$4
    local config_score=$5
    local cve_score=$6
    local total_score=$7
    local max_score=$8
    local percentage=$9
    local grade=${10}
    local minimalism_json=${11}
    local provenance_json=${12}
    local config_json=${13}
    local cve_json=${14}

    # Calculate individual section grades
    local -r minimalism_grade=$(get_grade "$minimalism_score" 4)
    local -r provenance_grade=$(get_grade "$provenance_score" 8)
    local -r config_grade=$(get_grade "$config_score" 4)
    local -r cve_grade=$(get_grade "$cve_score" 4)

    # Generate badge URLs
    local -r overall_badge=$(get_badge_url "overall" "$grade")
    local -r minimalism_badge=$(get_badge_url "minimalism" "$minimalism_grade")
    local -r provenance_badge=$(get_badge_url "provenance" "$provenance_grade")
    local -r config_badge=$(get_badge_url "configuration" "$config_grade")
    local -r cve_badge=$(get_badge_url "cves" "$cve_grade")

    cat << EOF
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>CHPs Scorer Report</title>
    <style>
        body { font-family: Arial, sans-serif; margin: 20px; }
        h1 { color: #333; }
        table { border-collapse: collapse; width: 100%; margin-top: 20px; }
        th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }
        th { background-color: #f4f4f4; }
        .badge { display: inline-block; margin: 5px 0; }
        .grade { font-weight: bold; }
    </style>
</head>
<body>
    <h1>CHPs Scorer Report</h1>
    <p><strong>Image:</strong> $image</p>
    <p><strong>Digest:</strong> $digest</p>
    <table>
        <thead>
            <tr>
                <th>Category</th>
                <th>Score</th>
                <th>Max Score</th>
                <th>Grade</th>
                <th>Badge</th>
            </tr>
        </thead>
        <tbody>
            <tr>
                <td>Minimalism</td>
                <td>$minimalism_score</td>
                <td>4</td>
                <td class="grade">$minimalism_grade</td>
                <td><img class="badge" src="$minimalism_badge" alt="Minimalism Badge"></td>
            </tr>
            <tr>
                <td>Provenance</td>
                <td>$provenance_score</td>
                <td>8</td>
                <td class="grade">$provenance_grade</td>
                <td><img class="badge" src="$provenance_badge" alt="Provenance Badge"></td>
            </tr>
            <tr>
                <td>Configuration</td>
                <td>$config_score</td>
                <td>4</td>
                <td class="grade">$config_grade</td>
                <td><img class="badge" src="$config_badge" alt="Configuration Badge"></td>
            </tr>
            <tr>
                <td>CVE</td>
                <td>$cve_score</td>
                <td>4</td>
                <td class="grade">$cve_grade</td>
                <td><img class="badge" src="$cve_badge" alt="CVE Badge"></td>
            </tr>
        </tbody>
        <tfoot>
            <tr>
                <th>Overall</th>
                <th>$total_score</th>
                <th>$max_score</th>
                <th class="grade">$grade</th>
                <th><img class="badge" src="$overall_badge" alt="Overall Badge"></th>
            </tr>
        </tfoot>
    </table>
    <p><strong>Overall Percentage:</strong> $percentage%</p>
</body>
</html>
EOF
}

# Function to output scores in JSON format
output_json() {
    local image=$1
    local digest=$2
    local minimalism_score=$3
    local provenance_score=$4
    local config_score=$5
    local cve_score=$6
    local total_score=$7
    local max_score=$8
    local percentage=$9
    local grade=${10}
    local minimalism_json=${11}
    local provenance_json=${12}
    local config_json=${13}
    local cve_json=${14}

    # Calculate individual section grades
    local -r minimalism_grade=$(get_grade "$minimalism_score" 4)
    local -r provenance_grade=$(get_grade "$provenance_score" 8)
    local -r config_grade=$(get_grade "$config_score" 4)
    local -r cve_grade=$(get_grade "$cve_score" 4)

    # Generate badge URLs
    local -r overall_badge=$(get_badge_url "overall" "$grade")
    local -r minimalism_badge=$(get_badge_url "minimalism" "$minimalism_grade")
    local -r provenance_badge=$(get_badge_url "provenance" "$provenance_grade")
    local -r config_badge=$(get_badge_url "configuration" "$config_grade")
    local -r cve_badge=$(get_badge_url "cves" "$cve_grade")

    cat << EOF
{
    "image": "$image",
    "digest": "$digest",
    "scores": {
        "minimalism": {
            "score": $minimalism_score,
            "max": 4,
            "grade": "$minimalism_grade",
            "badge": "$minimalism_badge",
            "checks": $(echo "$minimalism_json" | jq '.checks')
        },
        "provenance": {
            "score": $provenance_score,
            "max": 8,
            "grade": "$provenance_grade",
            "badge": "$provenance_badge",
            "checks": $(echo "$provenance_json" | jq '.checks')
        },
        "configuration": {
            "score": $config_score,
            "max": 4,
            "grade": "$config_grade",
            "badge": "$config_badge",
            "checks": $(echo "$config_json" | jq '.checks')
        },
        "cves": {
            "score": $cve_score,
            "max": 5,
            "grade": "$cve_grade",
            "badge": "$cve_badge",
            "checks": $(echo "$cve_json" | jq '.checks')
        }
    },
    "overall": {
        "score": $total_score,
        "max": $max_score,
        "percentage": $percentage,
        "grade": "$grade",
        "badge": "$overall_badge"
    }
}
EOF
}

# Function to output scores in text format
output_text() {
    local image=$1
    local digest=$2
    local minimalism_score=$3
    local provenance_score=$4
    local config_score=$5
    local cve_score=$6
    local total_score=$7
    local max_score=$8
    local percentage=$9
    local grade=${10}
    
    # Calculate individual section grades
    local -r minimalism_grade=$(get_grade "$minimalism_score" 4)
    local -r provenance_grade=$(get_grade "$provenance_score" 8)
    local -r config_grade=$(get_grade "$config_score" 4)
    local -r cve_grade=$(get_grade "$cve_score" 4)
    
    echo -e "${BOLD}Scoring image:${NC} $image"
    echo -e "${BOLD}Image digest:${NC} $digest"
    echo

    echo -e "${BOLD}Minimalism Score:${NC} $minimalism_score/4 $(get_grade_color "$minimalism_grade")($minimalism_grade)${NC}"

    echo -e "${BOLD}Provenance Score:${NC} $provenance_score/8 $(get_grade_color "$provenance_grade")($provenance_grade)${NC}"
    echo -e "${BOLD}Configuration Score:${NC} $config_score/4 $(get_grade_color "$config_grade")($config_grade)${NC}"
    echo -e "${BOLD}CVE Score:${NC} $cve_score/4 $(get_grade_color "$cve_grade")($cve_grade)${NC}"
    

    echo -e "${BOLD}Overall Score:${NC} $total_score/$max_score ($percentage%)"
    echo -e "${BOLD}Grade:${NC} $(get_grade_color "$grade")$grade${NC}"
    
    echo

    # Generate badge URLs
    local -r overall_badge=$(get_badge_url "overall" "$grade")
    local -r minimalism_badge=$(get_badge_url "minimalism" "$minimalism_grade")
    local -r provenance_badge=$(get_badge_url "provenance" "$provenance_grade")
    local -r config_badge=$(get_badge_url "configuration" "$config_grade")
    local -r cve_badge=$(get_badge_url "cves" "$cve_grade")

    # Check for terminal image support
    local term_support
    term_support=$(detect_term_img_support)
    local can_show_images=false
    
    if [[ "$term_support" != "none" ]] && check_curl; then
        can_show_images=true
    fi
    if [[ "$can_show_images" == "true" ]]; then
        display_badge "$minimalism_badge" "$term_support"
        display_badge "$provenance_badge" "$term_support"
        display_badge "$config_badge" "$term_support"
        display_badge "$cve_badge" "$term_support"
        display_badge "$overall_badge" "$term_support"
    else
        echo "![Minimalism Badge]($minimalism_badge)"
        echo "![Provenance Badge]($provenance_badge)"
        echo "![Configuration Badge]($config_badge)"
        echo "![CVE Badge]($cve_badge)"
        echo "![Overall Badge]($overall_badge)"
    fi
    
}

# Badges only
output_badges() {
    local image=$1
    local digest=$2
    local minimalism_score=$3
    local provenance_score=$4
    local config_score=$5
    local cve_score=$6
    local total_score=$7
    local max_score=$8
    local percentage=$9
    local grade=${10}

    # Generate badge URLs
    local -r overall_badge=$(get_badge_url "overall" "$grade")
    local -r minimalism_badge=$(get_badge_url "minimalism" "$minimalism_grade")
    local -r provenance_badge=$(get_badge_url "provenance" "$provenance_grade")
    local -r config_badge=$(get_badge_url "configuration" "$config_grade")
    local -r cve_badge=$(get_badge_url "cves" "$cve_grade")

    echo "![Minimalism Badge]($minimalism_badge)"
    echo "![Provenance Badge]($provenance_badge)"
    echo "![Configuration Badge]($config_badge)"
    echo "![CVE Badge]($cve_badge)"
    echo "![Overall Badge]($overall_badge)"
}

# Main scoring function
score_image() {
    local image=$1
    local dockerfile=$2
    local use_local_image=$3

    echo "Scoring image: $image" >&2
    if [ -n "$dockerfile" ]; then
        echo "Using Dockerfile: $dockerfile" >&2
    fi
    echo "----------------------------------------" >&2

    # Run minimalism checks
    minimalism_json=$(run_minimalism_checks "$image" "$dockerfile" "$use_local_image")
    minimalism_score=$(echo "$minimalism_json" | jq -r '.score')

    # Run provenance checks
    provenance_json=$(run_provenance_checks "$image" "$dockerfile" "$use_local_image")
    provenance_score=$(echo "$provenance_json" | jq -r '.score')

    # Run configuration checks
    config_json=$(run_config_checks "$image" "$dockerfile")
    config_score=$(echo "$config_json" | jq -r '.score')

    # Run CVE checks
    if [ "$SKIP_CVES" != "true" ]; then
        cve_json=$(run_cve_checks "$image")
        cve_score=$(echo "$cve_json" | jq -r '.score')
    else
        cve_score=0
        cve_json='{"score": 0, "checks": {}}'
    fi

    # Calculate overall score
    local total_score=$((minimalism_score + config_score + provenance_score + cve_score))
    local max_score=20  # Updated max score (4 + 4 + 8 + 4)
    local percentage=$((total_score * 100 / max_score))

    # Determine grade based on percentage
    local grade
    if [ $percentage -ge 94 ]; then  # 17-18 points
        grade="A+"
    elif [ $percentage -ge 75 ]; then  # 14-16 points
        grade="A"
    elif [ $percentage -ge 56 ]; then  # 11-13 points
        grade="B"
    elif [ $percentage -ge 38 ]; then  # 8-10 points
        grade="C"
    elif [ $percentage -ge 19 ]; then  # 5-7 points
        grade="D"
    else  # 0-4 points (Level None)
        grade="E"
    fi

    case "$OUTPUT_FORMAT" in
        json)
            output_json "$ORIGINAL_IMAGE" "$image" "$minimalism_score" "$provenance_score" "$config_score" "$cve_score" "$total_score" "$max_score" "$percentage" "$grade" "$minimalism_json" "$provenance_json" "$config_json" "$cve_json"
            ;;
        text)
            output_text "$ORIGINAL_IMAGE" "$image" "$minimalism_score" "$provenance_score" "$config_score" "$cve_score" "$total_score" "$max_score" "$percentage" "$grade"
            ;;
        html)
            output_html "$ORIGINAL_IMAGE" "$image" "$minimalism_score" "$provenance_score" "$config_score" "$cve_score" "$total_score" "$max_score" "$percentage" "$grade" "$minimalism_json" "$provenance_json" "$config_json" "$cve_json"
            ;;
        badges)
            output_badges "$ORIGINAL_IMAGE" "$image" "$minimalism_score" "$provenance_score" "$config_score" "$cve_score" "$total_score" "$max_score" "$percentage" "$grade"
            ;;
        *)
            echo "Error: Unknown output format '$OUTPUT_FORMAT'" >&2
            echo "Supported formats: text (default), json, html, badges" >&2
            exit 1
            ;;
    esac
}

# Check if Docker is running
if ! docker info >/dev/null 2>&1; then
    echo "Error: Docker is not running" >&2
    exit 1
fi

# Check if Dockerfile exists if provided
if [ -n "$DOCKERFILE" ] && [ ! -f "$DOCKERFILE" ]; then
    echo "Error: Dockerfile not found at $DOCKERFILE" >&2
    exit 1
fi


# Pull image or use a local one
#
# Note that local images do not have a repo digest.
if [ "$USE_LOCAL_IMAGE" == "true" ]; then
    echo "Using local image: $1"
    # local images do not have digest like remote, but instead they have Id
    if [ -n "$(docker inspect "$1" --format '{{.Id}}')" ]; then
        IMAGE_WITH_DIGEST=$(docker inspect "$1" --format '{{.Id}}')
    else
        exit 1
    fi
else
    echo "Pulling image: $1" >&2

    if ! docker pull "$1" > /dev/null 2>&1; then
        echo "Error: Failed to pull image" >&2
        exit 1
    fi

    # Get the full image name with digest
    IMAGE_WITH_DIGEST=$(docker inspect "$1" --format '{{.RepoDigests}}' 2>/dev/null | tr -d '[]' | cut -d' ' -f1)
fi

if [ -z "$IMAGE_WITH_DIGEST" ]; then
    echo "Warning: Could not get image digest, using original image name" >&2
    IMAGE_WITH_DIGEST="$1"
fi

ORIGINAL_IMAGE="$1"

# Run the scoring with the full image name including digest
score_image "$IMAGE_WITH_DIGEST" "$DOCKERFILE" "$USE_LOCAL_IMAGE"
