session.sh

#!/usr/bin/env bash
set -euo pipefail

# --------- Requirements / Config ----------
: "${S3_ENDPOINT_URL:?Set S3_ENDPOINT_URL (e.g., https://<accountid>.r2.cloudflarestorage.com)}"
: "${AWS_ACCESS_KEY_ID:?Set AWS_ACCESS_KEY_ID}"
: "${AWS_SECRET_ACCESS_KEY:?Set AWS_SECRET_ACCESS_KEY}"

INPUT_BUCKET="${INPUT_BUCKET:-inputs-bucket}"
OUTPUT_BUCKET="${OUTPUT_BUCKET:-outputs-bucket}"

WORKDIR="${WORKDIR:-$HOME/job}"
IN_DIR="$WORKDIR/in"
OUT_DIR="$WORKDIR/out"
LOG_DIR="$WORKDIR/logs"

mkdir -p "$IN_DIR" "$OUT_DIR" "$LOG_DIR"

# --------- Install s5cmd if missing ----------
need_s5cmd() {
  if command -v s5cmd >/dev/null 2>&1; then
    return 0
  fi
  sudo apt-get update -y
  sudo apt-get install -y wget ca-certificates
  local latest
  latest="$(wget -qO- https://api.github.com/repos/peak/s5cmd/releases/latest | grep -m1 '"tag_name"' | cut -d\" -f4)"
  wget -q "https://github.com/peak/s5cmd/releases/download/${latest}/s5cmd_${latest#v}_Linux-64bit.tar.gz" -O /tmp/s5cmd.tgz
  sudo tar -C /usr/local/bin -xzf /tmp/s5cmd.tgz s5cmd
  sudo chmod +x /usr/local/bin/s5cmd
}

S5="s5cmd --endpoint-url ${S3_ENDPOINT_URL}"

# --------- Helpers ----------
hr() { printf '%*s\n' "${COLUMNS:-80}" '' | tr ' ' -; }

# List objects under a prefix and return a numbered menu:
# We use `s5cmd ls s3://bucket/prefix` output.
list_prefix() {
  local bucket="$1"
  local prefix="$2"
  $S5 ls "s3://${bucket}/${prefix}" 2>/dev/null || true
}

# Parse ls output into:
# - folders (common prefixes ending with /)
# - files
# s5cmd output looks like:
#   2025/01/01 00:00:00     1234 some/key
#   DIR                       some/prefix/
browse_bucket() {
  local bucket="$1"
  local prefix="${2:-}"
  local selection=""
  local prompt=""

  while true; do
    clear || true
    echo "R2 Browser"
    hr
    echo "Bucket : ${bucket}"
    echo "Prefix : /${prefix}"
    hr

    local raw
    raw="$(list_prefix "$bucket" "$prefix")"
    if [[ -z "$raw" ]]; then
      echo "(No objects found under this prefix.)"
    else
      echo "$raw"
    fi

    hr
    echo "Commands:"
    echo "  cd <subprefix/>     - enter folder (must end with /)"
    echo "  up                  - go up one level"
    echo "  get <key>           - download file key to $IN_DIR"
    echo "  mget <pattern>      - download by wildcard (e.g. wordlists/*.zst)"
    echo "  pwd                 - show current prefix"
    echo "  done                - exit browser"
    hr
    read -r -p "> " prompt || true

    case "$prompt" in
      cd\ *)
        selection="${prompt#cd }"
        if [[ "$selection" != */ ]]; then
          echo "Subprefix must end with / (example: rules/). Press Enter."
          read -r
        else
          prefix="${prefix}${selection}"
        fi
        ;;
      up)
        # Remove trailing segment
        prefix="${prefix%/}"
        prefix="${prefix%/*}"
        [[ -n "$prefix" ]] && prefix="${prefix}/"
        ;;
      pwd)
        echo "/${prefix}"
        read -r -p "Press Enter..."
        ;;
      get\ *)
        selection="${prompt#get }"
        echo "Downloading s3://${bucket}/${prefix}${selection} -> ${IN_DIR}/"
        $S5 cp "s3://${bucket}/${prefix}${selection}" "$IN_DIR/"
        read -r -p "Done. Press Enter..."
        ;;
      mget\ *)
        selection="${prompt#mget }"
        echo "Downloading s3://${bucket}/${prefix}${selection} -> ${IN_DIR}/"
        $S5 cp "s3://${bucket}/${prefix}${selection}" "$IN_DIR/"
        read -r -p "Done. Press Enter..."
        ;;
      done)
        break
        ;;
      *)
        echo "Unknown command. Press Enter."
        read -r
        ;;
    esac
  done
}

upload_results() {
  local ts host out_prefix
  ts="$(date -u +%Y%m%dT%H%M%SZ)"
  host="$(hostname)"
  out_prefix="s3://${OUTPUT_BUCKET}/results/${host}/${ts}/"

  echo "Uploading ${OUT_DIR}/ -> ${out_prefix}"
  $S5 sync "${OUT_DIR}/" "${out_prefix}"
  echo "Upload complete."
}

main_menu() {
  while true; do
    clear || true
    echo "R2 Session Menu"
    hr
    echo "Workspace:"
    echo "  IN : $IN_DIR"
    echo "  OUT: $OUT_DIR"
    echo "  LOG: $LOG_DIR"
    hr
    echo "1) Browse INPUT bucket (${INPUT_BUCKET})"
    echo "2) Upload OUT results to OUTPUT bucket (${OUTPUT_BUCKET})"
    echo "3) Show environment status"
    echo "4) Exit"
    hr
    read -r -p "Select: " choice || true

    case "$choice" in
      1) browse_bucket "$INPUT_BUCKET" "" ;;
      2) upload_results; read -r -p "Press Enter..." ;;
      3)
        echo "S3_ENDPOINT_URL=$S3_ENDPOINT_URL"
        echo "INPUT_BUCKET=$INPUT_BUCKET"
        echo "OUTPUT_BUCKET=$OUTPUT_BUCKET"
        echo "WORKDIR=$WORKDIR"
        echo "s5cmd=$(command -v s5cmd || true)"
        read -r -p "Press Enter..."
        ;;
      4) exit 0 ;;
      *) echo "Invalid"; sleep 1 ;;
    esac
  done
}

need_s5cmd
main_menu