#!/bin/bash
# Copyright (c) StellaOps. All rights reserved.
# Licensed under the BUSL-1.1 license.
#
# collect-rekor-proofs.sh
# Collects Rekor transparency log inclusion proofs for release artifacts
#
# Usage: ./collect-rekor-proofs.sh --artifacts <dir> --output <dir>
#
# Prerequisites:
#   - rekor-cli installed (https://github.com/sigstore/rekor)
#   - Artifacts must already be signed and uploaded to Rekor

set -euo pipefail

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

# Default values
ARTIFACTS_DIR="artifacts"
OUTPUT_DIR="rekor-proofs"
REKOR_SERVER="${REKOR_SERVER:-https://rekor.sigstore.dev}"
PUBLIC_KEY_FILE="cosign.pub"

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --artifacts)
            ARTIFACTS_DIR="$2"
            shift 2
            ;;
        --output)
            OUTPUT_DIR="$2"
            shift 2
            ;;
        --public-key)
            PUBLIC_KEY_FILE="$2"
            shift 2
            ;;
        --rekor-server)
            REKOR_SERVER="$2"
            shift 2
            ;;
        --help)
            echo "Usage: $0 --artifacts <dir> --output <dir>"
            echo ""
            echo "Options:"
            echo "  --artifacts     Directory containing signed artifacts (default: artifacts)"
            echo "  --output        Output directory for Rekor proofs (default: rekor-proofs)"
            echo "  --public-key    Path to public key file (default: cosign.pub)"
            echo "  --rekor-server  Rekor server URL (default: https://rekor.sigstore.dev)"
            exit 0
            ;;
        *)
            echo -e "${RED}Unknown option: $1${NC}"
            exit 1
            ;;
    esac
done

# Check for rekor-cli
if ! command -v rekor-cli &> /dev/null; then
    echo -e "${YELLOW}Warning: rekor-cli not found. Skipping Rekor proof collection.${NC}"
    echo "Install from: https://github.com/sigstore/rekor/releases"
    mkdir -p "$OUTPUT_DIR"
    echo '{"warning": "rekor-cli not available", "proofs": []}' > "${OUTPUT_DIR}/inclusion-proofs.json"
    exit 0
fi

# Create output directories
mkdir -p "${OUTPUT_DIR}/log-entries"

echo -e "${GREEN}Collecting Rekor inclusion proofs${NC}"
echo "  Artifacts: ${ARTIFACTS_DIR}"
echo "  Output: ${OUTPUT_DIR}"
echo "  Rekor Server: ${REKOR_SERVER}"

# Initialize inclusion proofs JSON
proofs_json='{"proofs": []}'
checkpoint=""

# Function to collect proof for a single artifact
collect_proof() {
    local artifact_path="$1"
    local artifact_name
    artifact_name=$(basename "$artifact_path")
    local sig_path="${artifact_path}.sig"

    if [[ ! -f "$sig_path" ]]; then
        echo -e "  ${YELLOW}Skipping ${artifact_name}: no signature file found${NC}"
        return
    fi

    echo "  Processing: ${artifact_name}"

    # Search for the entry in Rekor
    local search_result
    if ! search_result=$(rekor-cli search --artifact "$artifact_path" --rekor_server "$REKOR_SERVER" 2>/dev/null); then
        echo -e "    ${YELLOW}No Rekor entry found${NC}"
        return
    fi

    # Extract UUIDs from search result
    local uuids
    uuids=$(echo "$search_result" | grep -oE '[0-9a-f]{64}' || true)

    if [[ -z "$uuids" ]]; then
        echo -e "    ${YELLOW}No matching entries in Rekor${NC}"
        return
    fi

    # Get the first (most recent) UUID
    local uuid
    uuid=$(echo "$uuids" | head -1)

    echo "    Found entry: ${uuid}"

    # Get the full log entry
    local entry_file="${OUTPUT_DIR}/log-entries/${uuid}.json"
    if rekor-cli get --uuid "$uuid" --rekor_server "$REKOR_SERVER" --format json > "$entry_file" 2>/dev/null; then
        echo -e "    ${GREEN}Saved log entry${NC}"

        # Extract log index and integrated time
        local log_index
        log_index=$(jq -r '.LogIndex' "$entry_file" 2>/dev/null || echo "-1")
        local integrated_time
        integrated_time=$(jq -r '.IntegratedTime' "$entry_file" 2>/dev/null || echo "0")

        # Add to proofs JSON
        proofs_json=$(echo "$proofs_json" | jq --arg uuid "$uuid" \
            --arg artifact "$artifact_name" \
            --argjson logIndex "$log_index" \
            --argjson integratedTime "$integrated_time" \
            --arg path "log-entries/${uuid}.json" \
            '.proofs += [{"uuid": $uuid, "artifactName": $artifact, "logIndex": $logIndex, "integratedTime": $integratedTime, "inclusionProofPath": $path}]')
    else
        echo -e "    ${YELLOW}Failed to retrieve entry details${NC}"
    fi
}

# Get Rekor checkpoint (signed tree head)
echo ""
echo "Fetching Rekor checkpoint..."
if checkpoint_result=$(curl -s "${REKOR_SERVER}/api/v1/log" 2>/dev/null); then
    echo "$checkpoint_result" > "${OUTPUT_DIR}/checkpoint.json"
    checkpoint=$(echo "$checkpoint_result" | jq -r '.signedTreeHead // empty' 2>/dev/null || true)
    if [[ -n "$checkpoint" ]]; then
        echo -e "  ${GREEN}Checkpoint saved${NC}"
    fi
fi

# Process all artifacts
echo ""
echo "Processing artifacts..."
for artifact in "${ARTIFACTS_DIR}"/stella-*.tar.gz "${ARTIFACTS_DIR}"/stella-*.zip; do
    if [[ -f "$artifact" ]]; then
        collect_proof "$artifact"
    fi
done

# Also process checksums if signed
for checksum_file in "${ARTIFACTS_DIR}"/*.sums "${ARTIFACTS_DIR}"/SHA256SUMS "${ARTIFACTS_DIR}"/SHA512SUMS; do
    if [[ -f "$checksum_file" ]] && [[ -f "${checksum_file}.sig" ]]; then
        collect_proof "$checksum_file"
    fi
done

# Write final inclusion proofs JSON
echo "$proofs_json" | jq '.' > "${OUTPUT_DIR}/inclusion-proofs.json"

# Count proofs
proof_count=$(echo "$proofs_json" | jq '.proofs | length')

echo ""
echo -e "${GREEN}Collected ${proof_count} inclusion proof(s)${NC}"
echo "Files written to: ${OUTPUT_DIR}/"
echo ""
echo "Contents:"
ls -la "${OUTPUT_DIR}/"