#! /bin/bash --
# sandbox-exec.sh — Run a command inside the sandbox (auto-selects backend)
#
# Usage:
#   sandbox-exec.sh [OPTIONS] -- CMD [ARGS...]
#
# Options:
#   --project-dir DIR         Directory with write access (default: $PWD)
#   --backend BACKEND         Force a specific backend (bwrap, landlock, auto)
#   --dry-run                 Print the sandbox command without executing
#   --cleanup-materialized    Remove BLOCKED_FILES placeholders we created
#                             at launch if they're still empty on exit.
#   --help                    Show this help
#
# Examples:
#   sandbox-exec.sh -- claude                        # Claude Code sandboxed
#   sandbox-exec.sh -- bash                          # interactive shell
#   sandbox-exec.sh --backend landlock -- bash       # force Landlock backend

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"

PROJECT_DIR=""
DRY_RUN=false
BACKEND_OVERRIDE=""
CLEANUP_MATERIALIZED_FLAG=""

VERSION="$(cat "$SCRIPT_DIR/VERSION" 2>/dev/null || echo "unknown")"

usage() {
    echo "agent-sandbox $VERSION — kernel-enforced isolation for AI coding agents"
    echo ""
    echo "Usage: sandbox-exec.sh [OPTIONS] -- CMD [ARGS...]"
    echo ""
    echo "Options:"
    echo "  --project-dir DIR         Directory with write access (default: \$PWD)"
    echo "  --backend BACKEND         Force backend: bwrap, landlock, auto (default: auto)"
    echo "  --dry-run                 Print the sandbox command without executing"
    echo "  --cleanup-materialized    Remove BLOCKED_FILES placeholders created at"
    echo "                            launch when they're still empty on sandbox exit."
    echo "                            (Also enabled by CLEANUP_MATERIALIZED_BLOCKED_FILES=1.)"
    echo "  --version                 Show version"
    echo "  --help                    Show this help"
    echo ""
    echo "Examples:"
    echo "  sandbox-exec.sh -- claude              # Claude Code in sandbox"
    echo "  sandbox-exec.sh -- bash                # interactive shell"
    echo "  sandbox-exec.sh --backend landlock -- cmd  # force Landlock"
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --project-dir)
            PROJECT_DIR="$2"
            shift 2
            ;;
        --backend)
            BACKEND_OVERRIDE="$2"
            shift 2
            ;;
        --dry-run)
            DRY_RUN=true
            shift
            ;;
        --cleanup-materialized)
            CLEANUP_MATERIALIZED_FLAG=1
            shift
            ;;
        --version)
            echo "agent-sandbox $VERSION"
            exit 0
            ;;
        --help|-h)
            usage
            exit 0
            ;;
        --)
            shift
            break
            ;;
        -*)
            echo "Error: Unknown option '$1'" >&2
            usage >&2
            exit 1
            ;;
        *)
            # No '--' before command — treat first non-option arg as the command
            break
            ;;
    esac
done

# Clear inherited sandbox env vars so Slurm jobs submitted from within a
# sandbox can start a fresh sandbox on the compute node.  This is safe
# because sandbox-exec.sh always sets up its own sandbox from scratch —
# the env vars are informational, not part of the enforcement mechanism
# (which is kernel-level: namespaces, Landlock rules, seccomp).
unset SANDBOX_ACTIVE SANDBOX_BACKEND SANDBOX_PROJECT_DIR

# Apply backend override before sourcing sandbox-lib.sh
if [[ -n "$BACKEND_OVERRIDE" ]]; then
    export SANDBOX_BACKEND="$BACKEND_OVERRIDE"
fi

# Save env overrides before config loading overwrites them.
# This allows users to override with: SLURM_SCOPE=session sandbox-exec.sh ...
_SLURM_SCOPE_ENV="${SLURM_SCOPE:-}"
_HOME_ACCESS_ENV="${HOME_ACCESS:-}"
_SANDBOX_QUIET_ENV="${SANDBOX_QUIET:-}"
_SANDBOX_NPROC_LIMIT_ENV="${SANDBOX_NPROC_LIMIT:-}"

source "$SCRIPT_DIR/sandbox-lib.sh"

# Default to current directory
if [[ -z "$PROJECT_DIR" ]]; then
    PROJECT_DIR="$(pwd)"
fi

# Resolve to absolute, physical path (follow symlinks).
# This ensures the project path inside the sandbox matches the path
# Claude Code sees outside, so session resume works across both.
PROJECT_DIR="$(cd "$PROJECT_DIR" && pwd -P)"

# Load per-project config overrides (conf.d/*.conf)
load_project_config "$PROJECT_DIR"

# Apply per-project environment overrides (SANDBOX_ENV from conf.d/*.conf).
# Done here so backends inherit the modified environment naturally —
# e.g. PATH modifications are picked up when bwrap/landlock/firejail
# prepend chaperon stubs.
for _env_entry in "${SANDBOX_ENV[@]}"; do
    export "$_env_entry"
done

# Validate
if [[ ! -d "$PROJECT_DIR" ]]; then
    echo "Error: Project directory does not exist: $PROJECT_DIR" >&2
    exit 1
fi

validate_project_dir "$PROJECT_DIR"

if [[ $# -eq 0 ]]; then
    echo "Error: No command specified after --" >&2
    echo "Hint: sandbox-exec.sh -- bash       (interactive shell)" >&2
    echo "      sandbox-exec.sh -- claude      (Claude Code)" >&2
    exit 1
fi

# Restore env overrides (saved before config loading)
if [[ -n "${_HOME_ACCESS_ENV:-}" ]]; then
    HOME_ACCESS="$_HOME_ACCESS_ENV"
fi
if [[ -n "${_SANDBOX_QUIET_ENV:-}" ]]; then
    SANDBOX_QUIET="$_SANDBOX_QUIET_ENV"
fi
if [[ -n "${_SANDBOX_NPROC_LIMIT_ENV:-}" ]]; then
    SANDBOX_NPROC_LIMIT="$_SANDBOX_NPROC_LIMIT_ENV"
fi

# Validate HOME_ACCESS
case "${HOME_ACCESS:-restricted}" in
    restricted|tmpwrite|read|write) ;;
    *) echo "Error: HOME_ACCESS must be 'restricted', 'tmpwrite', 'read', or 'write' (got '${HOME_ACCESS}')." >&2; exit 1 ;;
esac

# Load user-configured lmod modules (updates PATH for backend detection)
_load_sandbox_modules

# Detect and load backend
detect_backend
_BACKEND_DETECTED=true

# Export so the chaperon process (spawned below as a child) sees which
# backend it's running under. Used by chaperon/handlers/_handler_lib.sh
# to gate features that depend on backend capabilities — e.g. the Slurm
# --output / --error path transformation requires a RO-overlayable
# .sandbox-state/, which landlock cannot provide (additive-rules
# limitation), so the chaperon disables that feature on landlock.
export SANDBOX_BACKEND

# Apply per-agent grants. For each name in ENABLED_AGENTS, fold its
# declared writable/readable/blocked paths (from agents/<name>/config.conf)
# into the sandbox permission arrays. Must run after all config has
# been loaded so user-set ENABLED_AGENTS wins, and before the
# requirement check / overlay execution which both iterate the list.
_apply_agent_profiles

# Prepare agent profiles (backend-independent). Only agents listed in
# ENABLED_AGENTS are prepared. The requirement check emits warnings
# if declared credentials/paths look unreachable.
_check_agent_requirements
prepare_agent_configs "$PROJECT_DIR"

# Materialize BLOCKED_FILES placeholders, warn on each one we created,
# and track them so the EXIT trap can optionally clean up post-exit.
# Must run AFTER _apply_agent_profiles (so agent-added entries are
# folded in) and BEFORE backend_prepare (so backends see all targets
# as existing). See sandbox-lib.sh §_ensure_blocked_files_exist.
_ensure_blocked_files_exist

# Resolve --cleanup-materialized opt-in. Flag wins; env var is the
# config-file equivalent. Default off. Only meaningful when something
# was actually materialized this launch (the cleanup function is a
# no-op otherwise).
_CLEANUP_MATERIALIZED=0
if [[ "$CLEANUP_MATERIALIZED_FLAG" == "1" ]] \
   || _is_true "${CLEANUP_MATERIALIZED_BLOCKED_FILES:-false}"; then
    _CLEANUP_MATERIALIZED=1
fi

# ── Chaperon: create FIFO directory ───────────────────────────────
# Create the FIFO directory BEFORE backend_prepare so backends can
# add bind-mounts for it. The chaperon process is started AFTER
# backend_prepare but before entering the sandbox.

_CHAPERON_PID=""
_CHAPERON_DIR=""

# ── Proxy fallback: per-launch socket dir + daemon ────────────────
# Populated when `resolve_network_filter_mode` (called inside
# backend_prepare) resolves to `proxied`. Backends bind-mount
# $_NETWORK_PROXY_DIR into the sandbox; the proxy daemon spawns
# AFTER backend_prepare so the in-sandbox bind path is settled.
_NETWORK_PROXY_DIR=""
_NETWORK_PROXY_PID=""

# ── Mail-block: per-launch stub symlink-farm dir ──────────────────
# Populated when `resolve_network_mail_block_mode` (called inside
# backend_prepare) resolves to `on`. Holds one symlink per canonical
# mailer name pointing at the universal stub; the backend bind-mounts
# the dir at the same path on both sides of the bwrap boundary
# (mirrors the chaperon FIFO + proxy socket-dir pattern) and prepends
# that path to PATH inside the sandbox. Cleaned up on exit via the
# trap below.
_MAIL_BLOCK_STUBS_DIR=""

if [[ -x "$SCRIPT_DIR/chaperon/chaperon.sh" ]]; then
    _CHAPERON_DIR="$(mktemp -d "${TMPDIR:-/tmp}/chaperon-XXXXXX")"
    chmod 700 "$_CHAPERON_DIR"

    # request pipe: sandbox writes → chaperon reads
    mkfifo "$_CHAPERON_DIR/req"
    chmod 600 "$_CHAPERON_DIR/req"

    # Export the FIFO directory path — backends read this during prepare
    export _CHAPERON_FIFO_DIR="$_CHAPERON_DIR"
fi

# Prepare sandbox (reads _CHAPERON_FIFO_DIR for bind-mounts)
backend_prepare "$PROJECT_DIR"

if [[ "$DRY_RUN" == true ]]; then
    backend_dry_run "$@"
    # Clean up chaperon dir on dry-run
    [[ -n "$_CHAPERON_DIR" ]] && rm -rf "$_CHAPERON_DIR"
    exit 0
fi

# ── Proxy fallback: spawn host-side daemon (only when proxied) ────
#
# When backend_prepare resolved NETWORK_FILTER_MODE=proxied, it set
# _NETWORK_PROXY_DIR (the per-launch Unix-socket dir) and prepared the
# backend's bind-mount + env-var argv to expose it inside the sandbox.
# Spawn the host-side proxy server here, blocking until it writes
# "ready\n" on stdout (so the sandbox doesn't race past the socket
# bind).
#
# Lifecycle: PR_SET_PDEATHSIG (set inside the Python helper as its
# FIRST action) makes the proxy die when this shell dies, which is
# what happens when bwrap exits and the exec'd shell terminates.
# Belt-and-suspenders kill+rm-rf in _sandbox_cleanup covers the
# pre-exec failure window (Ctrl-C between proxy spawn and exec).
if [[ -n "${_NETWORK_PROXY_DIR:-}" ]]; then
    _proxy_py="$SCRIPT_DIR/tools/proxy/agent-sandbox-proxy.py"
    if [[ ! -r "$_proxy_py" ]]; then
        echo "Error: NETWORK_FILTER_MODE=proxied requested but $_proxy_py is missing." >&2
        exit 1
    fi
    # Compose the effective blocklist + exception list as JSON arrays
    # via Python's json.dumps so multi-line / quoted entries can't
    # shell-quote-escape into the argv. The shell pipes one entry per
    # line; Python collects them into a list.
    _proxy_blocklist_json="$(
        effective_network_blocklist | python3 -c 'import json,sys; print(json.dumps([l for l in sys.stdin.read().splitlines() if l]))'
    )"
    _proxy_except_json="$(
        effective_network_exception_list 2>/dev/null \
            | python3 -c 'import json,sys; print(json.dumps([l for l in sys.stdin.read().splitlines() if l]))'
    )"
    # Spawn the server. setsid detaches the daemon from the
    # controlling terminal so Ctrl-C in the parent shell goes to the
    # foreground bwrap, not the daemon — the daemon then dies via
    # PR_SET_PDEATHSIG when the parent exits.
    python3 "$_proxy_py" \
        --server --socket-dir "$_NETWORK_PROXY_DIR" \
        --blocklist-json "$_proxy_blocklist_json" \
        --except-json    "$_proxy_except_json" >"$_NETWORK_PROXY_DIR/ready" 2>"$_NETWORK_PROXY_DIR/proxy.err" &
    _NETWORK_PROXY_PID=$!
    # Wait ≤2s for the daemon to write "ready\n" to its stdout (we
    # tee'd that to a sentinel file). Python startup on cold-cache HPC
    # login nodes can take ~300ms; 2s is comfortable headroom.
    _ready_deadline=$(( SECONDS + 2 ))
    while (( SECONDS < _ready_deadline )); do
        if [[ -s "$_NETWORK_PROXY_DIR/ready" ]] \
           && grep -q '^ready$' "$_NETWORK_PROXY_DIR/ready" 2>/dev/null; then
            break
        fi
        sleep 0.02
    done
    if ! grep -q '^ready$' "$_NETWORK_PROXY_DIR/ready" 2>/dev/null; then
        echo "Error: agent-sandbox-proxy.py did not signal readiness within 2s." >&2
        [[ -s "$_NETWORK_PROXY_DIR/proxy.err" ]] && \
            sed 's/^/  proxy.err: /' "$_NETWORK_PROXY_DIR/proxy.err" >&2
        exit 1
    fi
fi

# Start chaperon in background.
# stdout → /dev/null so the chaperon doesn't hold the parent's output
# pipes open (otherwise $() subshells block until the chaperon exits).
# stderr → chaperon.err in the FIFO dir so crashes leave a trace.
# The chaperon's own per-request logging goes to structured log files;
# chaperon.err captures only unexpected failures (set -e, source errors,
# unhandled signals) that kill the process before logging is initialized.
if [[ -n "$_CHAPERON_DIR" ]]; then
    # Export config vars that chaperon handlers need (shell variables
    # are not inherited by child processes unless exported).
    # Restore env override if the user set SLURM_SCOPE before sandbox launch.
    if [[ -n "${_SLURM_SCOPE_ENV:-}" ]]; then
        SLURM_SCOPE="$_SLURM_SCOPE_ENV"
    fi
    export SLURM_SCOPE="${SLURM_SCOPE:-project}"
    export CHAPERON_LOG_LEVEL="${CHAPERON_LOG_LEVEL:-info}"
    export CHAPERON_LOG_RETAIN_DAYS="${CHAPERON_LOG_RETAIN_DAYS:-7}"

    "$SCRIPT_DIR/chaperon/chaperon.sh" \
        "$_CHAPERON_DIR" "$PROJECT_DIR" "$SCRIPT_DIR/sandbox-exec.sh" \
        >/dev/null 2>"$_CHAPERON_DIR/chaperon.err" &
    _CHAPERON_PID=$!
fi

# Clean up chaperon + proxy on exit — backend_exec does exec, so this
# only runs if exec fails (e.g., command not found, bwrap error). On
# success, exec replaces this shell and child processes die via
# PR_SET_PDEATHSIG (chaperon and the proxy server both arm it as their
# first action). Belt-and-suspenders: trap kills + cleanup.
_sandbox_cleanup() {
    if [[ -n "${_CHAPERON_PID:-}" ]]; then
        kill "$_CHAPERON_PID" 2>/dev/null || true
        wait "$_CHAPERON_PID" 2>/dev/null || true
    fi
    if [[ -n "${_CHAPERON_DIR:-}" ]]; then
        rm -rf "$_CHAPERON_DIR" 2>/dev/null || true
    fi
    if [[ -n "${_NETWORK_PROXY_PID:-}" ]]; then
        kill "$_NETWORK_PROXY_PID" 2>/dev/null || true
        wait "$_NETWORK_PROXY_PID" 2>/dev/null || true
    fi
    if [[ -n "${_NETWORK_PROXY_DIR:-}" ]]; then
        rm -rf "$_NETWORK_PROXY_DIR" 2>/dev/null || true
    fi
    if [[ -n "${_MAIL_BLOCK_STUBS_DIR:-}" ]]; then
        rm -rf "$_MAIL_BLOCK_STUBS_DIR" 2>/dev/null || true
    fi
    # Materialized BLOCKED_FILES cleanup. The function early-returns
    # when nothing was materialized this launch, so it's safe to call
    # unconditionally. Only reachable in the no-exec path (cleanup
    # mode); the exec'd path never runs the trap.
    if [[ "${_CLEANUP_MATERIALIZED:-0}" == "1" ]]; then
        _cleanup_materialized_blocked_files
    fi
}
trap _sandbox_cleanup EXIT

# Close inherited file descriptors (3+) to prevent FD-based sandbox bypass.
# Open FDs from the parent process persist across exec and bypass filesystem
# isolation — a parent-opened FD to /etc/shadow would remain readable inside
# the sandbox even if the path is blocked.  Keep only stdin/stdout/stderr.
# EXCEPTION: the chaperon socketpair FD must survive into the sandbox.
if [[ -d /proc/self/fd ]]; then
    for _fd in /proc/self/fd/*; do
        _fd_num="${_fd##*/}"
        if [[ "$_fd_num" -gt 2 ]] 2>/dev/null; then
            eval "exec ${_fd_num}>&-" 2>/dev/null || true
        fi
    done
fi

# Startup banner (disable with SANDBOX_QUIET=true in sandbox.conf or env)
if ! _is_true "${SANDBOX_QUIET:-false}"; then
    case "${HOME_ACCESS:-restricted}" in
        restricted) _home_label="restricted (listed dotfiles only, read-only)" ;;
        tmpwrite)   _home_label="tmpwrite (~ visible, writes to tmpfs — not persisted)" ;;
        read)       _home_label="read (~ visible, read-only)" ;;
        write)      _home_label="write (~ visible, read-write)" ;;
        *)          _home_label="${HOME_ACCESS}" ;;
    esac
    echo "sandbox: $SANDBOX_BACKEND | project: $PROJECT_DIR | home: $_home_label" >&2
fi

backend_exec "$@"
