#!/bin/bash
set -euo pipefail

# Portable root detection - works anywhere
ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
mkdir -p "$LOGS"

cd "$ROOT"

LIBHEIF_DIR="$ROOT/external/libheif"
POC_SRC="$ROOT/repro/poc.c"
POC_BIN="$ROOT/repro/poc"
POC_FILE="$ROOT/repro/poc.heif"
MANIFEST="$ROOT/repro/runtime_manifest.json"

# Build vulnerable version if needed
if [ ! -f "$LIBHEIF_DIR/build_vuln/libheif/libheif.so" ]; then
    echo "[BUILD] Building vulnerable libheif v1.21.2 with ASan..."
    cd "$LIBHEIF_DIR"
    git checkout v1.21.2
    mkdir -p build_vuln
    cd build_vuln
    cmake .. -DWITH_UNCOMPRESSED_CODEC=ON -DWITH_JPEG_ENCODER=ON -DWITH_JPEG_DECODER=ON \
        -DCMAKE_BUILD_TYPE=Debug \
        -DCMAKE_C_FLAGS="-fsanitize=address -g" -DCMAKE_CXX_FLAGS="-fsanitize=address -g" \
        -DWITH_EXAMPLES=OFF
    make -j$(nproc) heif
fi

# Build fixed version if needed
if [ ! -f "$LIBHEIF_DIR/build_fix/libheif/libheif.so" ]; then
    echo "[BUILD] Building fixed libheif v1.22.0 with ASan..."
    cd "$LIBHEIF_DIR"
    git checkout v1.22.0
    mkdir -p build_fix
    cd build_fix
    cmake .. -DWITH_UNCOMPRESSED_CODEC=ON -DWITH_JPEG_ENCODER=ON -DWITH_JPEG_DECODER=ON \
        -DCMAKE_BUILD_TYPE=Debug \
        -DCMAKE_C_FLAGS="-fsanitize=address -g" -DCMAKE_CXX_FLAGS="-fsanitize=address -g" \
        -DWITH_EXAMPLES=OFF
    make -j$(nproc) heif
fi

# Compile PoC harness against a given build directory
compile_poc() {
    local build_dir="$1"
    echo "[COMPILE] Compiling PoC against $build_dir"
    gcc -fsanitize=address -g "$POC_SRC" \
        -I"$LIBHEIF_DIR/libheif/api" -I"$build_dir" \
        -L"$build_dir/libheif" \
        -lheif -Wl,-rpath,"$build_dir/libheif" \
        -o "$POC_BIN"
}

# Generate PoC file (same file works for both builds)
if [ ! -f "$POC_FILE" ]; then
    compile_poc "$LIBHEIF_DIR/build_vuln"
    "$POC_BIN" "$POC_FILE"
fi

# --- Vulnerable run ---
compile_poc "$LIBHEIF_DIR/build_vuln"
echo "[TEST] Running PoC against vulnerable libheif v1.21.2..."
ASAN_OPTIONS=detect_leaks=0 "$POC_BIN" "$POC_FILE" decode > "$LOGS/vuln_stdout.txt" 2> "$LOGS/vuln_stderr.txt" || true

VULN_HAS_ASAN=0
if grep -q "ERROR: AddressSanitizer: heap-buffer-overflow" "$LOGS/vuln_stderr.txt"; then
    VULN_HAS_ASAN=1
fi

# --- Fixed run ---
compile_poc "$LIBHEIF_DIR/build_fix"
echo "[TEST] Running PoC against fixed libheif v1.22.0..."
ASAN_OPTIONS=detect_leaks=0 "$POC_BIN" "$POC_FILE" decode > "$LOGS/fix_stdout.txt" 2> "$LOGS/fix_stderr.txt" || true

FIX_HAS_ASAN=0
if grep -q "ERROR: AddressSanitizer:" "$LOGS/fix_stderr.txt"; then
    FIX_HAS_ASAN=1
fi

# --- Write runtime manifest ---
cat > "$MANIFEST" <<EOF
{
  "poc_file": "$POC_FILE",
  "poc_sha256": "$(sha256sum "$POC_FILE" | awk '{print $1}')",
  "vulnerable_build": "v1.21.2",
  "fixed_build": "v1.22.0",
  "vulnerable_asan_detected": $VULN_HAS_ASAN,
  "fixed_asan_detected": $FIX_HAS_ASAN,
  "vulnerable_stderr": "$LOGS/vuln_stderr.txt",
  "fixed_stderr": "$LOGS/fix_stderr.txt",
  "timestamp": "$(date -u +%Y-%m-%dT%H:%M:%SZ)"
}
EOF

# --- Verdict ---
if [ "$VULN_HAS_ASAN" -eq 1 ] && [ "$FIX_HAS_ASAN" -eq 0 ]; then
    echo "[RESULT] CONFIRMED: ASan heap-buffer-overflow in v1.21.2, clean in v1.22.0"
    exit 0
else
    echo "[RESULT] UNEXPECTED: vuln_asan=$VULN_HAS_ASAN fixed_asan=$FIX_HAS_ASAN"
    exit 1
fi