#!/bin/bash
set -e

DIR=/root/w/packet_ebpf
CERTS=/tmp/ebpf_certs
PIDS=()

# Cleanup on exit
cleanup() { kill "${PIDS[@]}" 2>/dev/null; rm -rf "$CERTS"; }
trap cleanup EXIT

# Start background process and track PID
bg() { "$@" &>/dev/null & PIDS+=($!); }

# Print result
ok()   { echo "  $1: OK"; }
fail() { echo "  $1: FAIL"; }

# Build
echo "=== Building ==="
cd "$DIR"
cargo xtask build-ebpf --release &>/dev/null
cargo build --release -p packet-detector &>/dev/null
echo "Done"
echo

# Test 1: UDP - test actual packet delivery
echo "=== UDP Magic Word (only 'hell0123' passes) ==="

# Start XDP filter
./target/release/packet-detector lo &>/dev/null & PIDS+=($!)
sleep 1

RECV_FILE=$(mktemp)

# Test valid packet
timeout 2 bash -c "nc -u -l 127.0.0.1 9999 > $RECV_FILE" &
sleep 0.3
echo 'hell0123' | nc -u -w1 127.0.0.1 9999 2>/dev/null || true
sleep 0.5
grep -q 'hell0123' "$RECV_FILE" && ok "hell0123 passed (XDP_PASS)" || fail "hell0123 dropped"

# Test invalid packet
> "$RECV_FILE"  # clear file
timeout 2 bash -c "nc -u -l 127.0.0.1 9999 > $RECV_FILE" &
sleep 0.3
echo 'wrongmsg' | nc -u -w1 127.0.0.1 9999 2>/dev/null || true
sleep 0.5
grep -q 'wrongmsg' "$RECV_FILE" && fail "wrongmsg passed (should drop)" || ok "wrongmsg dropped (XDP_DROP)"

rm -f "$RECV_FILE"
kill "${PIDS[-1]}" 2>/dev/null; unset 'PIDS[-1]'
echo

# Test 2: TLS
echo "=== TLS Certificate ==="

# Create two separate PKI environments
BAD_CERTS="$CERTS/bad"
GOOD_CERTS="$CERTS/good"
mkdir -p "$BAD_CERTS" "$GOOD_CERTS"

# Start server with UNTRUSTED certs on port 8443
cd "$BAD_CERTS"
bg "$DIR/target/release/tls_server" 8443
sleep 3

# Start server with TRUSTED certs on port 8444
cd "$GOOD_CERTS"
bg "$DIR/target/release/tls_server" 8444
sleep 3

# Start packet-detector with ONLY the good CA (won't trust bad server)
"$DIR/target/release/packet-detector" lo "$GOOD_CERTS/ca_cert.pem" &>/dev/null & PIDS+=($!)
sleep 2

# Test 1: Connect to BAD server (untrusted cert)
# Handshake completes, but HTTP request should fail (blocked after validation)
cd "$BAD_CERTS"
BAD_OUT=$(timeout 5 "$DIR/target/release/tls_client" 127.0.0.1 8443 2>&1) || true
if echo "$BAD_OUT" | grep -q "HTTP/1.1 200"; then
    fail "Untrusted cert - HTTP should have been blocked"
else
    ok "Untrusted cert - HTTP blocked (XDP_DROP after validation)"
fi

# Test 2: Connect to GOOD server (trusted cert) - should succeed fully
cd "$GOOD_CERTS"
GOOD_OUT=$(timeout 5 "$DIR/target/release/tls_client" 127.0.0.1 8444 2>&1) || true
if echo "$GOOD_OUT" | grep -q "HTTP/1.1 200"; then
    ok "Trusted cert - HTTP succeeded (XDP_PASS)"
else
    fail "Trusted cert - HTTP failed"
fi

echo
echo "=== Done ==="