#!/bin/sh
# Stands up a real WireGuard tunnel entirely inside this machine, so the
# server's production network posture -- "bind wg0 and nothing else, fail
# closed if it's missing" -- can be exercised without a phone, a router
# port-forward, or any internet exposure.
#
# The shape, all in one kernel:
#
#     main netns                        "phone" netns
#     wg0  10.66.0.1  <-- encrypted --> wg1  10.66.0.2
#       |                                 |
#     veth-srv 10.99.0.1 <--- UDP ---> veth-phone 10.99.0.2
#
# The veth pair stands in for "the internet" carrying WireGuard's UDP; the
# wg interfaces are real, with a real handshake and real keys. 10.66.0.1 is
# deliberately the same address the leaf certificate carries a SAN for
# (certs.rs covers every local address), so a client inside the tunnel
# completes the same pinned-TLS handshake a phone will.
#
#   scripts/test-wg-tunnel.sh up      create the tunnel (needs sudo)
#   scripts/test-wg-tunnel.sh test    run the server on wg0 and reach it from "phone"
#   scripts/test-wg-tunnel.sh down    remove everything it created
#
# Everything here is torn down by `down`: the netns (taking wg1 and the veth
# peer with it), wg0, and the temporary key files.

set -eu

NS=phone
WG_SERVER=wg0
WG_CLIENT=wg1
SERVER_WG_IP=10.66.0.1
CLIENT_WG_IP=10.66.0.2
SERVER_UDP_IP=10.99.0.1
CLIENT_UDP_IP=10.99.0.2
LISTEN_PORT=51820
KEYDIR=/run/ai-app-wg-test
REPO=$(cd "$(dirname "$0")/.." && pwd)

up() {
    echo "==> Generating ephemeral keypairs in $KEYDIR"
    sudo mkdir -p "$KEYDIR"
    sudo sh -c "umask 077; wg genkey > $KEYDIR/server.key; wg genkey > $KEYDIR/client.key"
    sudo sh -c "wg pubkey < $KEYDIR/server.key > $KEYDIR/server.pub"
    sudo sh -c "wg pubkey < $KEYDIR/client.key > $KEYDIR/client.pub"

    echo "==> Creating netns '$NS' and the veth pair that carries the UDP"
    sudo ip netns add "$NS"
    sudo ip link add veth-srv type veth peer name veth-phone
    sudo ip link set veth-phone netns "$NS"
    sudo ip addr add "$SERVER_UDP_IP/24" dev veth-srv
    sudo ip link set veth-srv up
    sudo ip -n "$NS" addr add "$CLIENT_UDP_IP/24" dev veth-phone
    sudo ip -n "$NS" link set veth-phone up
    sudo ip -n "$NS" link set lo up

    echo "==> Creating $WG_SERVER (server side, $SERVER_WG_IP)"
    sudo ip link add "$WG_SERVER" type wireguard
    sudo sh -c "wg set $WG_SERVER listen-port $LISTEN_PORT private-key $KEYDIR/server.key \
        peer \$(cat $KEYDIR/client.pub) allowed-ips $CLIENT_WG_IP/32"
    sudo ip addr add "$SERVER_WG_IP/24" dev "$WG_SERVER"
    sudo ip link set "$WG_SERVER" up

    # Moving the interface preserves its birth namespace's UDP socket.
    echo "==> Creating $WG_CLIENT (phone side, $CLIENT_WG_IP) in netns '$NS'"
    sudo ip link add "$WG_CLIENT" type wireguard
    sudo ip link set "$WG_CLIENT" netns "$NS"
    sudo ip netns exec "$NS" sh -c "wg set $WG_CLIENT private-key $KEYDIR/client.key \
        peer \$(cat $KEYDIR/server.pub) allowed-ips $SERVER_WG_IP/32 \
        endpoint $SERVER_UDP_IP:$LISTEN_PORT persistent-keepalive 5"
    sudo ip -n "$NS" addr add "$CLIENT_WG_IP/24" dev "$WG_CLIENT"
    sudo ip -n "$NS" link set "$WG_CLIENT" up

    echo "==> Forcing a handshake"
    sudo ip netns exec "$NS" ping -c 2 -W 3 "$SERVER_WG_IP" >/dev/null 2>&1 || true
    sudo wg show "$WG_SERVER" | sed 's/^/    /'
    echo "==> Up. wg0 is $SERVER_WG_IP; run '$0 test' next."
}

test_tunnel() {
    CERTS="${XDG_CONFIG_HOME:-$HOME/.config}/ai-app/certs"
    if [ ! -f "$CERTS/leaf.pem" ]; then
        echo "No certificates in $CERTS -- start ai-server once; it makes them." >&2
        exit 1
    fi
    if [ ! -x "$REPO/server/target/debug/ai-server" ]; then
        echo "Build the server first: (cd server && cargo build)" >&2
        exit 1
    fi

    echo "==> Starting ai-server with NO --bind (production path: wg0 only)"
    setsid nohup "$REPO/server/target/debug/ai-server" \
        </dev/null >"$REPO/server/wg-test.log" 2>&1 &
    sleep 2

    echo "==> Where is it actually listening?"
    ss -tlnp 2>/dev/null | grep 8443 | sed 's/^/    /' || echo "    (nothing on 8443)"

    echo "==> From inside the tunnel: GET /sessions through wg1 -> wg0"
    if [ -z "${AI_TOKEN:-}" ]; then
        echo "    (set AI_TOKEN=<the enrollment token> to test an authorized call;"
        echo "     without it this only proves reachability + TLS, via a 401)"
    fi
    sudo ip netns exec "$NS" curl -s -o /dev/null -w "    HTTP %{http_code} (TLS ok, pinned CA)\n" \
        --cacert "$CERTS/ca.pem" \
        ${AI_TOKEN:+-H "Authorization: Bearer $AI_TOKEN"} \
        "https://$SERVER_WG_IP:8443/sessions" || echo "    UNREACHABLE"

    echo "==> Handshake counters (proves the traffic really crossed WireGuard)"
    sudo wg show "$WG_SERVER" transfer | sed 's/^/    /'

    pkill -f "[a]i-server" || true
    echo "==> Server stopped."
}

down() {
    echo "==> Removing tunnel"
    sudo ip netns del "$NS" 2>/dev/null || true
    sudo ip link del "$WG_SERVER" 2>/dev/null || true
    sudo ip link del veth-srv 2>/dev/null || true
    sudo rm -rf "$KEYDIR"
    echo "==> Down."
}

case "${1:-}" in
    up) up ;;
    test) test_tunnel ;;
    down) down ;;
    *) echo "usage: $0 up|test|down" >&2; exit 1 ;;
esac