#!/bin/bash
# Apply the DNS jail to this Explore container, and install `unjail` / `rejail`.
#
# Explore is meant to behave like a trial: the session captured here becomes the trial's
# seed, so an agent that reached the network here would produce a snapshot the trial
# cannot reproduce. Same jail, applied every boot (docker remounts /etc/resolv.conf per
# start, so it cannot be baked into the image).
#
# Live resolution only — no address pinning. An Explore container can run for days, so a
# resolved-at-boot address has far longer to go stale than in a single trial.
set -u

JAIL_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
STATE=/tmp/.dnsjail

[ "${RACCOON_DNS_JAIL:-0}" = "1" ] || exit 0

# Only the model endpoint gates the jail. The toolkit's telemetry hosts go in as extras
# (below): those sends are backgrounded and disowned, so one failing to resolve would fail
# silently rather than visibly -- and must not take the whole jail down with it.
allow_hosts() {
    local url="${ANTHROPIC_BASE_URL:-}" host=""
    [ -n "$url" ] || return 1
    host="${url#*://}"; host="${host%%/*}"; host="${host##*@}"; host="${host%%:*}"
    [ -n "$host" ] || return 1
    case "$host" in *[!A-Za-z0-9.-]* | -* | .* | *.) return 1 ;; esac
    printf '%s' "$host"
}

install_helpers() {
    sudo tee /usr/local/bin/unjail >/dev/null <<'EOF'
#!/bin/sh
# Restore this container's DNS. The jail comes back on the next container start, or now
# with `rejail`. Package installs need this; run-app does it for you around its own.
[ -f /tmp/.dnsjail/resolv.orig ] || { echo "unjail: not jailed"; exit 0; }
sudo sh -c 'cat /tmp/.dnsjail/resolv.orig > /etc/resolv.conf'
echo "unjail: DNS restored — run 'rejail' when you are done, or restart the container."
EOF
    sudo tee /usr/local/bin/rejail >/dev/null <<EOF
#!/bin/sh
[ -f /tmp/.dnsjail/allow ] || { echo "rejail: nothing to restore"; exit 1; }
sudo env DNSJAIL_ALLOW="\$(cat /tmp/.dnsjail/allow)" \
    DNSJAIL_ALLOW_EXTRA="\$(cat /tmp/.dnsjail/allow-extra 2>/dev/null)" \
    sh $JAIL_DIR/dns-jail-container.sh
grep -qE '^nameserver[[:space:]]+127\.0\.0\.1[[:space:]]*$' /etc/resolv.conf && echo "rejail: jailed" || echo "rejail: could not jail — left as is"
EOF
    sudo chmod +x /usr/local/bin/unjail /usr/local/bin/rejail
}

# Not fatal: an Explore container that cannot jail is still a usable Explore container.
dnsjail_off() {
    mkdir -p "$STATE" 2>/dev/null || true
    printf '%s\n' "$1" > "$STATE/why" 2>/dev/null || true
    echo "dns-jail: off for this session — normal network access. Not an error."
    exit 0
}

[ -f "$JAIL_DIR/dns-jail-container.sh" ] || dnsjail_off "script not present: $JAIL_DIR/dns-jail-container.sh"
# Jailing without the model endpoint on the allowlist would strand the agent, so a
# missing or unusable ANTHROPIC_BASE_URL means no jail at all.
ALLOW="$(allow_hosts)" || dnsjail_off "no usable host in ANTHROPIC_BASE_URL: ${ANTHROPIC_BASE_URL:-<unset>}"
# Parent domains for the telemetry, not the exact endpoints: both CNAME within their own
# domain, and the catch-all would NXDOMAIN a chain target that is not itself allowed.
sudo env DNSJAIL_ALLOW="$ALLOW" \
    DNSJAIL_ALLOW_EXTRA="amplitude.com datadoghq.com ${RACCOON_DNS_JAIL_ALLOW:-}" \
    sh "$JAIL_DIR/dns-jail-container.sh" || true
install_helpers

# Report what the script decided, rather than re-probing: it already verified the model
# endpoint against its own resolver and failed open if that did not hold. A second probe
# here has to pick a control host -- and any host the worker allowlists makes that control
# resolve, reading a working jail as a broken one and tearing it down.
if grep -qE '^nameserver[[:space:]]+127\.0\.0\.1[[:space:]]*$' /etc/resolv.conf; then
    echo "dns-jail: DNS limited to the model endpoint and toolkit telemetry."
    echo "          Installing packages? \`unjail\` (then \`rejail\`). run-app handles its own."
else
    dnsjail_off "the jail did not take; see $STATE/dnsmasq.err if present"
fi