#!/usr/bin/env bash
# sandbox.sh - Safe AI coding agent sandbox
# Usage: ./sandbox.sh <command> [options]
# Commands: setup, run, shell, status, clean
set -euo pipefail

# ── Config ────────────────────────────────────────────────────────────────────
SANDBOX_IMAGE="${SANDBOX_IMAGE:-ai-sandbox:latest}"
SANDBOX_DIR="${SANDBOX_DIR:-$HOME/.ai-sandbox}"
WORKDIR_HOST="${SANDBOX_DIR}/workspace"
CONFIG_DIR="${SANDBOX_DIR}/config"      # persistent agent credentials (host-only, chmod 700)
TMPDIR_SIZE="${TMPDIR_SIZE:-512m}"
CPU_LIMIT="${CPU_LIMIT:-1.0}"
MEM_LIMIT="${MEM_LIMIT:-512m}"
NETWORK_MODE="${NETWORK_MODE:-none}"   # none = no internet; bridge = allow

# Colors
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; CYAN='\033[0;36m'; NC='\033[0m'
info()    { echo -e "${CYAN}[sandbox]${NC} $*"; }
ok()      { echo -e "${GREEN}[ok]${NC} $*"; }
warn()    { echo -e "${YELLOW}[warn]${NC} $*"; }
die()     { echo -e "${RED}[error]${NC} $*" >&2; exit 1; }

# ── OS detection ──────────────────────────────────────────────────────────────
OS="$(uname -s)"
case "$OS" in
  Linux)  PLATFORM=linux ;;
  Darwin) PLATFORM=macos ;;
  *)      die "Unsupported OS: $OS" ;;
esac

# ── Helpers ───────────────────────────────────────────────────────────────────
require_docker() {
  command -v docker &>/dev/null || die "Docker not found. Run: ./sandbox.sh setup"
  docker info &>/dev/null       || die "Docker daemon not running."
}

has_nsjail() {
  [[ "$PLATFORM" == "linux" ]] && command -v nsjail &>/dev/null
}

has_firejail() {
  command -v firejail &>/dev/null
}

# Verify that host networking/DNS was not changed by this script
assert_host_untouched() {
  # We never modify /etc/hosts, /etc/resolv.conf, iptables, or network interfaces.
  # This function just documents that guarantee and does a quick sanity check.
  if [[ "$PLATFORM" == "linux" ]]; then
    local resolv_hash_before resolv_hash_after
    resolv_hash_before="${_RESOLV_HASH_BEFORE:-}"
    resolv_hash_after="$(md5sum /etc/resolv.conf 2>/dev/null | cut -d' ' -f1)"
    if [[ -n "$resolv_hash_before" && "$resolv_hash_before" != "$resolv_hash_after" ]]; then
      die "/etc/resolv.conf was modified unexpectedly! Aborting."
    fi
  fi
}

snapshot_host() {
  if [[ "$PLATFORM" == "linux" ]]; then
    export _RESOLV_HASH_BEFORE="$(md5sum /etc/resolv.conf 2>/dev/null | cut -d' ' -f1)"
  fi
}

# ── Dockerfile (embedded) ─────────────────────────────────────────────────────
write_dockerfile() {
  local dir="$1"
  cat > "$dir/Dockerfile" <<'EOF'
FROM ubuntu:24.04

ARG DEBIAN_FRONTEND=noninteractive

# System packages
RUN apt-get update -qq && apt-get install -y --no-install-recommends \
    python3 python3-pip python3-venv \
    git curl wget ca-certificates gnupg \
    build-essential \
    && apt-get clean && rm -rf /var/lib/apt/lists/*

# Node.js LTS (via NodeSource — newer than Ubuntu's bundled version)
RUN curl -fsSL https://deb.nodesource.com/setup_lts.x | bash - \
    && apt-get install -y --no-install-recommends nodejs \
    && apt-get clean && rm -rf /var/lib/apt/lists/*

# GitHub CLI
RUN curl -fsSL https://cli.github.com/packages/githubcli-archive-keyring.gpg \
      | gpg --dearmor -o /usr/share/keyrings/githubcli-archive-keyring.gpg \
    && echo "deb [arch=$(dpkg --print-architecture) \
         signed-by=/usr/share/keyrings/githubcli-archive-keyring.gpg] \
         https://cli.github.com/packages stable main" \
      > /etc/apt/sources.list.d/github-cli.list \
    && apt-get update -qq && apt-get install -y --no-install-recommends gh \
    && apt-get clean && rm -rf /var/lib/apt/lists/*

# AI CLI tools (npm global)
RUN npm install -g --no-audit --no-fund \
    @anthropic-ai/claude-code \
    @google/gemini-cli

# Non-root sandbox user (fixed uid/gid 999)
RUN groupadd -r -g 999 agent \
    && useradd -r -u 999 -g agent -m -d /home/agent agent

# Workspace dir
RUN mkdir -p /workspace && chown agent:agent /workspace

USER agent
ENV HOME=/home/agent

# Stub dirs for bind/tmpfs mounts (must exist before runtime mounts)
RUN mkdir -p \
    /home/agent/.claude \
    /home/agent/.cache \
    /home/agent/.config/gh \
    /home/agent/.config/gemini \
    /home/agent/.local/state

# gh Copilot extension — baked into image; read-only at runtime is fine
RUN gh extension install github/gh-copilot 2>/dev/null || true

USER root
WORKDIR /workspace
USER agent

ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1

CMD ["/bin/bash"]
EOF
}

# ── Commands ──────────────────────────────────────────────────────────────────

cmd_setup() {
  snapshot_host
  info "Setting up AI sandbox on $PLATFORM..."

  # 1. Docker
  if ! command -v docker &>/dev/null; then
    if [[ "$PLATFORM" == "linux" ]]; then
      info "Installing Docker (rootless)..."
      # Use official convenience script but only if user confirms
      read -rp "  Install Docker via get.docker.com? [y/N] " ans
      [[ "$ans" =~ ^[Yy]$ ]] || die "Aborted. Install Docker manually: https://docs.docker.com/engine/install/"
      curl -fsSL https://get.docker.com | sh
      # Rootless setup (optional, avoids root daemon)
      if command -v dockerd-rootless-setuptool.sh &>/dev/null; then
        info "Configuring rootless Docker..."
        dockerd-rootless-setuptool.sh install 2>/dev/null || warn "Rootless setup skipped."
      fi
    elif [[ "$PLATFORM" == "macos" ]]; then
      if command -v brew &>/dev/null; then
        info "Installing Docker via Homebrew..."
        brew install --cask docker
        info "Please open Docker.app to finish installation, then re-run: ./sandbox.sh setup"
        open /Applications/Docker.app 2>/dev/null || true
        exit 0
      else
        die "Install Docker Desktop from https://docs.docker.com/desktop/mac/install/"
      fi
    fi
  else
    ok "Docker already installed: $(docker --version)"
  fi

  # 2. Verify Docker daemon
  if ! docker info &>/dev/null; then
    die "Docker daemon not running. Start Docker and retry."
  fi

  # 3. Optional: nsjail on Linux
  if [[ "$PLATFORM" == "linux" ]] && ! command -v nsjail &>/dev/null; then
    warn "nsjail not found (optional). For extra hardening: sudo apt install nsjail"
  elif has_nsjail; then
    ok "nsjail available: $(nsjail --version 2>&1 | head -1)"
  fi

  # 4. Create sandbox dirs
  mkdir -p "$WORKDIR_HOST"
  ok "Workspace: $WORKDIR_HOST"

  # Persistent config dirs — host-only, never world-readable
  mkdir -p \
    "$CONFIG_DIR/claude" \
    "$CONFIG_DIR/gh" \
    "$CONFIG_DIR/gemini"
  chmod 700 "$CONFIG_DIR" "$CONFIG_DIR/claude" "$CONFIG_DIR/gh" "$CONFIG_DIR/gemini"
  ok "Config dir:  $CONFIG_DIR (chmod 700)"

  # 5. Build Docker image
  local build_dir
  build_dir="$(mktemp -d)"
  write_dockerfile "$build_dir"
  info "Building sandbox image ($SANDBOX_IMAGE)..."
  docker build -t "$SANDBOX_IMAGE" "$build_dir" --quiet
  rm -rf "$build_dir"
  ok "Image built: $SANDBOX_IMAGE"

  assert_host_untouched
  echo ""
  ok "Setup complete. Host network/DNS: unchanged."
  info "Run agent code:  ./sandbox.sh run <script.py>"
  info "Interactive:     ./sandbox.sh shell"
}

cmd_run() {
  require_docker
  snapshot_host

  local script="${1:-}"
  [[ -z "$script" ]] && die "Usage: ./sandbox.sh run <script> [args...]"
  [[ -f "$script" ]] || die "File not found: $script"

  local abs_script
  abs_script="$(realpath "$script")"
  local script_dir
  script_dir="$(dirname "$abs_script")"
  local script_file
  script_file="$(basename "$abs_script")"

  info "Running in sandbox: $script_file"
  info "  Network:  $NETWORK_MODE"
  info "  CPU:      $CPU_LIMIT"
  info "  Memory:   $MEM_LIMIT"
  [[ "$NETWORK_MODE" == "none" ]] && info "  Internet: BLOCKED"

  # Detect language
  local runner=""
  case "$script_file" in
    *.py)   runner="python3" ;;
    *.js)   runner="node" ;;
    *.sh)   runner="bash" ;;
    *)      runner="" ;;  # executable or unknown
  esac

  local docker_args=(
    docker run
    --rm
    --network "$NETWORK_MODE"
    --cpus "$CPU_LIMIT"
    --memory "$MEM_LIMIT"
    --memory-swap "$MEM_LIMIT"                    # no swap
    --read-only                                   # read-only rootfs
    --tmpfs /tmp:size=128m,noexec                 # writable tmp, no exec
    --tmpfs /home/agent/.cache:size=128m,uid=999,gid=999
    --tmpfs /home/agent/.local/state:size=64m,uid=999,gid=999
    --cap-drop ALL                                # drop all Linux capabilities
    --security-opt no-new-privileges              # no privilege escalation
    --security-opt seccomp=unconfined
    --pids-limit 128                              # no fork bombs
    -v "${script_dir}:/workspace:ro"              # mount script dir read-only
    -v "${CONFIG_DIR}/claude:/home/agent/.claude:rw"
    -v "${CONFIG_DIR}/gh:/home/agent/.config/gh:rw"
    -v "${CONFIG_DIR}/gemini:/home/agent/.config/gemini:rw"
    --workdir /workspace
    "$SANDBOX_IMAGE"
  )

  # Optional nsjail wrapper on Linux
  if has_nsjail && [[ "$PLATFORM" == "linux" ]]; then
    info "  Extra:    nsjail wrapping enabled"
    # nsjail wraps the entire docker run call for an extra namespace layer
    # Note: this requires nsjail in PATH and appropriate permissions
    exec nsjail \
      --mode o \
      --quiet \
      --rlimit_nofile 64 \
      --rlimit_nproc 32 \
      --time_limit 300 \
      --disable_clone_newcgroup \
      -- "${docker_args[@]}" ${runner:+$runner} "$script_file" "${@:2}"
  fi

  "${docker_args[@]}" ${runner:+$runner} "$script_file" "${@:2}"

  assert_host_untouched
}

cmd_shell() {
  require_docker
  snapshot_host

  info "Opening interactive sandbox shell..."
  info "  Network: $NETWORK_MODE"
  [[ "$NETWORK_MODE" == "none" ]] && warn "  Internet is BLOCKED (set NETWORK_MODE=bridge to allow)"

  mkdir -p "$WORKDIR_HOST"
  mkdir -p "$CONFIG_DIR/claude" "$CONFIG_DIR/gh" "$CONFIG_DIR/gemini"

  docker run \
    --rm -it \
    --network "$NETWORK_MODE" \
    --cpus "$CPU_LIMIT" \
    --memory "$MEM_LIMIT" \
    --memory-swap "$MEM_LIMIT" \
    --read-only \
    --tmpfs /tmp:size=128m,noexec \
    --tmpfs /home/agent/.cache:size=128m,uid=999,gid=999 \
    --tmpfs /home/agent/.local/state:size=64m,uid=999,gid=999 \
    --tmpfs /workspace:size=256m,uid=999,gid=999 \
    --cap-drop ALL \
    --security-opt no-new-privileges \
    --pids-limit 128 \
    -v "${WORKDIR_HOST}:/mnt/workspace:rw" \
    -v "${CONFIG_DIR}/claude:/home/agent/.claude:rw" \
    -v "${CONFIG_DIR}/gh:/home/agent/.config/gh:rw" \
    -v "${CONFIG_DIR}/gemini:/home/agent/.config/gemini:rw" \
    "$SANDBOX_IMAGE" \
    /bin/bash

  assert_host_untouched
}

cmd_status() {
  require_docker
  info "Sandbox image:"
  docker images "$SANDBOX_IMAGE" 2>/dev/null || warn "Image not built yet."
  echo ""
  info "Running sandbox containers:"
  docker ps --filter "ancestor=$SANDBOX_IMAGE" --format "table {{.ID}}\t{{.Status}}\t{{.CreatedAt}}" 2>/dev/null || true
  echo ""
  info "Workspace: $WORKDIR_HOST"
  echo ""
  info "Optional tools:"
  has_nsjail   && ok "nsjail:   available" || warn "nsjail:   not found (optional)"
  has_firejail && ok "firejail: available" || warn "firejail: not found (optional)"
}

cmd_clean() {
  require_docker
  info "Removing sandbox image and workspace..."
  read -rp "  Remove image '$SANDBOX_IMAGE' and $WORKDIR_HOST? [y/N] " ans
  [[ "$ans" =~ ^[Yy]$ ]] || { info "Aborted."; exit 0; }
  docker rmi "$SANDBOX_IMAGE" 2>/dev/null && ok "Image removed." || warn "Image not found."
  rm -rf "$WORKDIR_HOST" && ok "Workspace removed."
}

cmd_help() {
  cat <<EOF

  ${CYAN}sandbox.sh${NC} - Safe AI coding agent sandbox

  ${YELLOW}Commands:${NC}
    setup              Check/install deps, build sandbox image
    run <script> [..] Run a script inside the sandbox (network OFF by default)
    shell              Interactive bash shell inside the sandbox
    status             Show image, running containers, available tools
    clean              Remove image and workspace
    help               This help

  ${YELLOW}Installed AI agents:${NC}
    claude               Claude Code CLI  (@anthropic-ai/claude-code)
    gemini               Gemini CLI       (@google/gemini-cli)
    gh copilot           GitHub Copilot   (gh extension)
    gh                   GitHub CLI       (for auth + copilot)

  ${YELLOW}Persistent credentials:${NC}
    ~/.ai-sandbox/config/claude/   → /home/agent/.claude
    ~/.ai-sandbox/config/gh/       → /home/agent/.config/gh
    ~/.ai-sandbox/config/gemini/   → /home/agent/.config/gemini
    (host dirs are chmod 700; survive container restarts)

  ${YELLOW}Environment:${NC}
    NETWORK_MODE=none    (default) No internet inside sandbox
    NETWORK_MODE=bridge  Allow internet (use with caution)
    MEM_LIMIT=512m       Memory cap (default: 512m)
    CPU_LIMIT=1.0        CPU cores (default: 1.0)
    SANDBOX_IMAGE=...    Docker image name (default: ai-sandbox:latest)

  ${YELLOW}Examples:${NC}
    ./sandbox.sh setup
    ./sandbox.sh run agent.py
    ./sandbox.sh shell
    NETWORK_MODE=bridge ./sandbox.sh run fetch_data.py

  ${YELLOW}Platforms:${NC}
    Linux:  Docker + optional nsjail hardening
    macOS:  Docker only

EOF
}

# ── Entrypoint ────────────────────────────────────────────────────────────────
CMD="${1:-help}"
shift || true

case "$CMD" in
  setup)   cmd_setup "$@" ;;
  run)     cmd_run "$@" ;;
  shell)   cmd_shell "$@" ;;
  status)  cmd_status "$@" ;;
  clean)   cmd_clean "$@" ;;
  help|-h|--help) cmd_help ;;
  *) die "Unknown command: $CMD. Run ./sandbox.sh help" ;;
esac