From c86dca283f59ec75ea593fca6410780b52ddf458 Mon Sep 17 00:00:00 2001
From: Justin Oros <JustinOros@gmail.com>
Date: Thu, 16 Apr 2026 09:30:47 -0700
Subject: [PATCH] add retry or abort prompt to all connection tests

---
 hub/onboard-spoke.sh | 36 ++++++++++++++++++++++++++----------
 setup.sh             | 27 ++++++++++++++++++++++-----
 2 files changed, 48 insertions(+), 15 deletions(-)

diff --git a/hub/onboard-spoke.sh b/hub/onboard-spoke.sh
index 2085639..826b79f 100644
--- a/hub/onboard-spoke.sh
+++ b/hub/onboard-spoke.sh
@@ -19,6 +19,25 @@ YELLOW='\033[1;33m'
 CYAN='\033[0;36m'
 NC='\033[0m'
 
+retry_or_abort() {
+    local test_cmd="$1"
+    local fail_msg="$2"
+    while true; do
+        if eval "$test_cmd" 2>/dev/null; then
+            return 0
+        fi
+        echo ""
+        warn "$fail_msg"
+        echo -e "  ${YELLOW}[R]${NC} Retry   ${RED}[A]${NC} Abort"
+        read -rp "Choice: " CHOICE
+        case "${CHOICE,,}" in
+            r) info "Retrying..." ;;
+            a) die "Aborted." ;;
+            *) warn "Press R to retry or A to abort." ;;
+        esac
+    done
+}
+
 info()   { echo -e "${GREEN}[+]${NC} $*"; }
 warn()   { echo -e "${YELLOW}[!]${NC} $*"; }
 die()    { echo -e "${RED}[ERROR]${NC} $*" >&2; exit 1; }
@@ -45,11 +64,9 @@ KEYSCAN=$(ssh-keyscan -p "$TUNNEL_PORT" -H localhost 2>/dev/null)
 echo "$KEYSCAN" >> "$SSH_DIR/known_hosts"
 
 info "Verifying spoke is reachable on port $TUNNEL_PORT..."
-if ssh -o BatchMode=yes -o ConnectTimeout=10 -p "$TUNNEL_PORT" armbian@localhost exit 2>/dev/null; then
-    info "Spoke is reachable."
-else
-    warn "Could not verify spoke without a key yet — proceeding to key setup."
-fi
+retry_or_abort \
+    "ssh -o BatchMode=yes -o ConnectTimeout=10 -p \"$TUNNEL_PORT\" armbian@localhost exit" \
+    "Spoke not reachable on port $TUNNEL_PORT. Make sure the tunnel is up."
 
 header "Generating Hub SSH Key"
 if [ -f "$KEY_PATH" ]; then
@@ -66,11 +83,10 @@ ssh-copy-id -i "$KEY_PATH.pub" -p "$TUNNEL_PORT" armbian@localhost
 info "Key copied."
 
 header "Testing Hub -> Spoke Key Auth"
-if ssh -i "$KEY_PATH" -o BatchMode=yes -o ConnectTimeout=10 -p "$TUNNEL_PORT" armbian@localhost exit 2>/dev/null; then
-    info "Key auth to spoke successful."
-else
-    die "Key auth failed. Check authorized_keys on the spoke."
-fi
+retry_or_abort \
+    "ssh -i \"$KEY_PATH\" -o BatchMode=yes -o ConnectTimeout=10 -p \"$TUNNEL_PORT\" armbian@localhost exit" \
+    "Key auth failed. Check authorized_keys on the spoke."
+info "Key auth to spoke successful."
 
 header "Adding rclone Remote"
 if grep -q "\[${SPOKE_NAME}-remote\]" "$RCLONE_CONF" 2>/dev/null; then
diff --git a/setup.sh b/setup.sh
index ca3bebc..3846ce9 100644
--- a/setup.sh
+++ b/setup.sh
@@ -16,6 +16,25 @@ YELLOW='\033[1;33m'
 CYAN='\033[0;36m'
 NC='\033[0m'
 
+retry_or_abort() {
+    local test_cmd="$1"
+    local fail_msg="$2"
+    while true; do
+        if eval "$test_cmd" 2>/dev/null; then
+            return 0
+        fi
+        echo ""
+        warn "$fail_msg"
+        echo -e "  ${YELLOW}[R]${NC} Retry   ${RED}[A]${NC} Abort"
+        read -rp "Choice: " CHOICE
+        case "${CHOICE,,}" in
+            r) info "Retrying..." ;;
+            a) die "Aborted." ;;
+            *) warn "Press R to retry or A to abort." ;;
+        esac
+    done
+}
+
 info()   { echo -e "${GREEN}[+]${NC} $*"; }
 warn()   { echo -e "${YELLOW}[!]${NC} $*"; }
 die()    { echo -e "${RED}[ERROR]${NC} $*" >&2; exit 1; }
@@ -107,11 +126,9 @@ sudo -u armbian ssh-keyscan -H "$HUB_HOST" >> "$SSH_DIR/known_hosts" 2>/dev/null
 
 header "Testing SSH Connection"
 info "Testing connection to $HUB_HOST..."
-if sudo -u armbian ssh -i "$KEY_PATH" -o BatchMode=yes -o ConnectTimeout=10 "$HUB_USER@$HUB_HOST" exit 2>/dev/null; then
-    info "SSH connection successful."
-else
-    die "SSH connection to $HUB_HOST failed. Check that finn added your public key."
-fi
+retry_or_abort \
+    "sudo -u armbian ssh -i \"$KEY_PATH\" -o BatchMode=yes -o ConnectTimeout=10 \"$HUB_USER@$HUB_HOST\" exit" \
+    "SSH connection to $HUB_HOST failed. Check that finn added your public key."
 
 header "Finding Available Tunnel Port"
 info "Scanning for a free port on $HUB_HOST starting from $START_PORT..."