#!/bin/bash
set -euo pipefail

ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
REPO="$ROOT/mongoose"
PATCH_FILE="$ROOT/coding/proposed_fix.diff"
WORK_DIR="$ROOT/coding/tmp_verify"
LOG_DIR="$ROOT/coding/logs"
VULN_COMMIT="0d882f1b^"

mkdir -p "$WORK_DIR" "$LOG_DIR"

if [ ! -d "$REPO/.git" ]; then
  echo "[!] Repository not found: $REPO"
  exit 1
fi

if [ ! -f "$PATCH_FILE" ]; then
  echo "[!] Patch file not found: $PATCH_FILE"
  exit 1
fi

SERVER_PID=""
cleanup() {
  if [ -n "$SERVER_PID" ] && kill -0 "$SERVER_PID" 2>/dev/null; then
    kill "$SERVER_PID" 2>/dev/null || true
    wait "$SERVER_PID" 2>/dev/null || true
  fi
}
trap cleanup EXIT

assert_contains() {
  local file="$1"
  local needle="$2"
  local reason="$3"
  if ! grep -Fq "$needle" "$file"; then
    echo "[!] Assertion failed: $reason"
    echo "[!] Expected to find: $needle"
    echo "[!] In file: $file"
    tail -n 120 "$file" || true
    exit 1
  fi
}

assert_not_contains() {
  local file="$1"
  local needle="$2"
  local reason="$3"
  if grep -Fq "$needle" "$file"; then
    echo "[!] Assertion failed: $reason"
    echo "[!] Unexpected content: $needle"
    echo "[!] In file: $file"
    tail -n 120 "$file" || true
    exit 1
  fi
}

cat > "$WORK_DIR/mdns_verify_harness.c" << 'EOF'
#include <stdio.h>
#include "mongoose.h"

#define MDNS_NAME "test"

static 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;
  (void) c;

  if (ev == MG_EV_MDNS_REQ && req != NULL && req->rr != NULL) {
    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);
      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\n");
      fflush(stdout);
    }
  }
}

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) {
    fprintf(stderr, "failed to start mdns listener\n");
    return 1;
  }

  printf("[*] mDNS server ready on UDP/5353\n");
  fflush(stdout);

  for (int i = 0; i < 240; i++) mg_mgr_poll(&mgr, 50);

  mg_mgr_free(&mgr);
  return 0;
}
EOF

cat > "$WORK_DIR/send_ptr_query.py" << 'EOF'
#!/usr/bin/env python3
import socket
import struct
import time

service = "_http._tcp.local"
header = struct.pack(">HHHHHH", 0x1234, 0x0000, 1, 0, 0, 0)
labels = service.split(".")
qname = b"".join(bytes([len(label)]) + label.encode() for label in labels) + b"\x00"
question = qname + struct.pack(">HH", 12, 1)  # PTR IN
packet = header + question

targets = ["224.0.0.251", "127.0.0.1"]

for i in range(15):
    for target in targets:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 1)
        sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_LOOP, 1)
        sock.sendto(packet, (target, 5353))
        sock.close()
    time.sleep(0.15)

print("sent PTR queries to multicast and loopback")
EOF
chmod +x "$WORK_DIR/send_ptr_query.py"

build_binary() {
  local tag="$1"
  gcc -g -fsanitize=address -fno-omit-frame-pointer -O0 \
    -DMG_ENABLE_MDNS=1 -DMG_ENABLE_LINES \
    -I"$REPO" \
    "$REPO/mongoose.c" "$WORK_DIR/mdns_verify_harness.c" \
    -o "$WORK_DIR/mdns_verify" > "$LOG_DIR/build_${tag}.log" 2>&1
}

run_case() {
  local tag="$1"
  local server_log="$LOG_DIR/${tag}_server.log"
  local trigger_log="$LOG_DIR/${tag}_trigger.log"

  "$WORK_DIR/mdns_verify" > "$server_log" 2>&1 &
  SERVER_PID=$!
  sleep 2

  python3 "$WORK_DIR/send_ptr_query.py" > "$trigger_log" 2>&1 || true
  sleep 4

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

echo "[+] Resetting to vulnerable commit: $VULN_COMMIT"
git -C "$REPO" reset --hard "$VULN_COMMIT" > /dev/null
git -C "$REPO" clean -fd > /dev/null

echo "[+] Building vulnerable target"
build_binary "vuln"

echo "[+] Running vulnerable target (expect overflow)"
run_case "vuln"
assert_contains "$LOG_DIR/vuln_server.log" "AddressSanitizer: stack-buffer-overflow" \
  "vulnerable build should trigger stack overflow"

echo "[+] Applying fix patch"
git -C "$REPO" apply "$PATCH_FILE"

echo "[+] Building patched target"
build_binary "fixed"

echo "[+] Running patched target (expect no overflow)"
run_case "fixed"
assert_not_contains "$LOG_DIR/fixed_server.log" "AddressSanitizer: stack-buffer-overflow" \
  "patched build should not trigger stack overflow"
assert_contains "$LOG_DIR/fixed_server.log" "[+] Got mDNS request! Type=12" \
  "patched build should still process incoming PTR request"

echo "[+] Verification successful"
echo "    Vulnerable run: overflow reproduced"
echo "    Patched run: overflow blocked"