|
| 1 | +#!/bin/bash |
| 2 | + |
| 3 | +start_nats_server() { |
| 4 | + # Run nats-server and purge its logs |
| 5 | + nats-server --config /etc/nats/nats-server.conf > /dev/null 2>&1 & |
| 6 | + |
| 7 | + # Wait for nats-server to be ready |
| 8 | + timeout_seconds=10 |
| 9 | + start_time=$(date +%s.%N) |
| 10 | + deadline_time=$(echo "$start_time + $timeout_seconds" | bc -l) |
| 11 | + while ! timeout 1 nats --user=sys --password=syspass server ping >/dev/null 2>&1; do |
| 12 | + sleep 0.2 |
| 13 | + current_time=$(date +%s.%N) |
| 14 | + if (( $(echo "$current_time > $deadline_time" | bc -l) )); then |
| 15 | + echo "⏰ Nats server failed to start after ${timeout_seconds} seconds." |
| 16 | + exit 1 |
| 17 | + fi |
| 18 | + done |
| 19 | +} |
| 20 | + |
| 21 | +stop_nats_server() { |
| 22 | + # kill nats-server |
| 23 | + kill -9 $(pgrep nats-server) |
| 24 | + |
| 25 | + # wait for nats-server to exit |
| 26 | + wait $(pgrep nats-server) |
| 27 | +} |
| 28 | + |
| 29 | +parse_args() { |
| 30 | + TIMEOUT=10 |
| 31 | + while [[ "$#" -gt 0 ]]; do |
| 32 | + case $1 in |
| 33 | + --timeout) |
| 34 | + TIMEOUT="$2" |
| 35 | + shift 2 |
| 36 | + ;; |
| 37 | + *) |
| 38 | + COMMAND_ARGS=("$@") |
| 39 | + break |
| 40 | + ;; |
| 41 | + esac |
| 42 | + done |
| 43 | +} |
| 44 | + |
| 45 | +read_command_from_stdin() { |
| 46 | + if [ ${#COMMAND_ARGS[@]} -eq 0 ] && [ ! -t 0 ]; then |
| 47 | + COMMAND=$(cat -) |
| 48 | + COMMAND_ARGS=($COMMAND) |
| 49 | + fi |
| 50 | +} |
| 51 | + |
| 52 | +preprocess_command() { |
| 53 | + if [ ${#COMMAND_ARGS[@]} -eq 0 ]; then |
| 54 | + echo "No command provided." |
| 55 | + exit 1 |
| 56 | + fi |
| 57 | + if [ ${#COMMAND_ARGS[@]} -eq 1 ] && [[ "${COMMAND_ARGS[0]}" == *" "* ]]; then |
| 58 | + COMMAND_ARGS=(${COMMAND_ARGS[0]}) |
| 59 | + fi |
| 60 | + COMMAND_STR="${COMMAND_ARGS[*]}" |
| 61 | + |
| 62 | + # Handle line continuations |
| 63 | + COMMAND_STR=$(echo "$COMMAND_STR" | sed 's/\\[[:space:]]/ /g') |
| 64 | + COMMAND_STR=$(echo "$COMMAND_STR" | sed 's/\\$//g') |
| 65 | +} |
| 66 | + |
| 67 | +add_auth_flags() { |
| 68 | + # Replace both && and ; with a single delimiter for splitting |
| 69 | + # IFS cannot handle multi-character sequences like &&, so we use sed to replace them |
| 70 | + SPLIT_CLEANED_COMMAND=$(echo "$COMMAND_STR" | sed 's/&&/|/g; s/;/|/g') |
| 71 | + |
| 72 | + IFS='|' read -ra SUBCOMMANDS <<< "$SPLIT_CLEANED_COMMAND" |
| 73 | + AUTH_COMMANDS=() |
| 74 | + for subcommand in "${SUBCOMMANDS[@]}"; do |
| 75 | + subcommand=$(echo "$subcommand" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//') |
| 76 | + if [[ -z "$subcommand" ]]; then |
| 77 | + continue |
| 78 | + fi |
| 79 | + if [[ "$subcommand" =~ ^nats ]]; then |
| 80 | + # Remove any existing auth flags |
| 81 | + subcommand=$(echo "$subcommand" | sed -E 's/--user(=[^ ]+| [^ ]+)? ?//g; s/--password(=[^ ]+| [^ ]+)? ?//g') |
| 82 | + |
| 83 | + # Add auth flags |
| 84 | + auth_subcommand=$(echo "$subcommand" | sed 's/^nats/& --user=ruser --password=T0pS3cr3t/') |
| 85 | + AUTH_COMMANDS+=("$auth_subcommand") |
| 86 | + else |
| 87 | + AUTH_COMMANDS+=("$subcommand") |
| 88 | + fi |
| 89 | + done |
| 90 | + PROCESSED_COMMAND=$(printf '%s && ' "${AUTH_COMMANDS[@]}" | sed 's/ && $//') |
| 91 | +} |
| 92 | + |
| 93 | +run_commands() { |
| 94 | + start_nats_server |
| 95 | + timeout "$TIMEOUT" bash -c "$PROCESSED_COMMAND" |
| 96 | + if [ $? -eq 124 ]; then |
| 97 | + stop_nats_server |
| 98 | + echo "⏰ Execution timed out after ${TIMEOUT} seconds." |
| 99 | + exit 124 |
| 100 | + fi |
| 101 | +} |
| 102 | + |
| 103 | +main() { |
| 104 | + parse_args "$@" |
| 105 | + read_command_from_stdin |
| 106 | + preprocess_command |
| 107 | + add_auth_flags |
| 108 | + run_commands |
| 109 | +} |
| 110 | + |
| 111 | +main "$@" |
0 commit comments