From 82b865391606a695148b488a44078a3d38de0f00 Mon Sep 17 00:00:00 2001
From: lilyLuLiu <lilyliulu777@outlook.com>
Date: Thu, 24 Oct 2024 11:48:05 +0800
Subject: [PATCH] multiple attempts for scp and ssh command fix failed ssh
 connection when using password

---
 entrypoint.sh        | 15 +++++++++------
 lib/common/remote.sh | 45 +++++++++++++++++++++++++++++++++++---------
 2 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/entrypoint.sh b/entrypoint.sh
index f9ca5c6..e8041fa 100755
--- a/entrypoint.sh
+++ b/entrypoint.sh
@@ -21,7 +21,7 @@ if [[ ! remote_required ]] || [[ ! mamp_required ]] || [[ -z "${ASSETS_FOLDER+x}
 fi
 
 # Create ssh config file is using proxy connection
-if [[ -n "${BASTION_HOST}" && -n "${BASTION_HOST_USERNAME}" ]];then
+if [[ -n "${BASTION_HOST+x}" && -n "${BASTION_HOST_USERNAME+x}" ]];then
     ssh_config_file
 fi
 
@@ -37,14 +37,16 @@ fi
 # Create execution folder 
 echo "Create assets folder on target"
 TARGET_FOLDER="${TARGET_FOLDER:-"deliverest-${RANDOM}"}"
-$(ssh_cmd "mkdir -p ${TARGET_FOLDER}")
+mkdir_cmd="$(ssh_cmd mkdir -p ${TARGET_FOLDER})"
+exec_and_retry 5 5 ${mkdir_cmd}
 
 # Copy asset
 echo "Copy assets folder to target"
-$(scp_to_cmd "${ASSETS_FOLDER}/*" "${TARGET_FOLDER}/")
+scp_cmd="$(scp_to_cmd "${ASSETS_FOLDER}/*" "${TARGET_FOLDER}/")"
+exec_and_retry 5 5 ${scp_cmd}
 
 # Exec command
-$(ssh_cmd "$@")
+exec_and_retry 5 5 "$(ssh_cmd $@)"
 
 # If remote workload includes a reboot this is the only way to ensure we can 
 # copy results if any or cleanup
@@ -60,11 +62,12 @@ fi
 if [[ ! -z "${TARGET_RESULTS+x}" ]]; then
     # If exec create some reuslts we define the env and they will be copied to OUTPUT_FOLDER
     OUTPUT_FOLDER="${OUTPUT_FOLDER:-"/output"}"
-    $(scp_from_cmd "${TARGET_FOLDER}/${TARGET_RESULTS}" "${OUTPUT_FOLDER}/")
+    scp_cmd="$(scp_from_cmd "${TARGET_FOLDER}/${TARGET_RESULTS}" "${OUTPUT_FOLDER}/)")"
+    exec_and_retry 5 5 ${scp_cmd}
 fi
 
 if [ "${TARGET_CLEANUP:-}" = "true" ]; then
     # This will create the cmd based on OS env with the right syntax
     cmd="$(remove_folder ${TARGET_FOLDER})"
-    $(ssh_cmd ${cmd})
+    exec_and_retry 5 5 "$(ssh_cmd $cmd)"
 fi
\ No newline at end of file
diff --git a/lib/common/remote.sh b/lib/common/remote.sh
index 060779e..a92a9dd 100755
--- a/lib/common/remote.sh
+++ b/lib/common/remote.sh
@@ -27,7 +27,9 @@ connect_options() {
     options="$options -o UserKnownHostsFile=/dev/null"
     options="$options -o ServerAliveInterval=30"
     options="$options -o ServerAliveCountMax=1200"
-    options="$options -o BatchMode=yes"
+    if [[ -n "${TARGET_HOST_KEY_PATH+x}" ]]; then
+        options="$options -o BatchMode=yes"
+    fi
     options="$options -o ConnectTimeout=3"
     echo $options
 }
@@ -85,7 +87,7 @@ check_connection() {
 # Define remote connection
 uri () {
     local remote="${TARGET_HOST_USERNAME}@${TARGET_HOST}"
-    if [[ -n "${TARGET_HOST_DOMAIN}" ]]; then
+    if [[ -n "${TARGET_HOST_DOMAIN+x}" ]]; then
         remote="${TARGET_HOST_USERNAME}@${TARGET_HOST_DOMAIN}@${TARGET_HOST}"
     fi
     echo "${remote}" 
@@ -96,9 +98,9 @@ uri () {
 # $2 remote path
 scp_to_cmd () {
     cmd="scp -r $(connect_options) "
-    if [[ -n "${BASTION_HOST}" && -n "${BASTION_HOST_USERNAME}" ]]; then
+    if [[ -n "${BASTION_HOST+x}" && -n "${BASTION_HOST_USERNAME+x}" ]]; then
         echo "${cmd} -F ssh_config ${1} target_host:${2}"
-    elif [[ -n "${TARGET_HOST_KEY_PATH}" ]]; then
+    elif [[ -n "${TARGET_HOST_KEY_PATH+x}" ]]; then
         echo "${cmd} -i ${TARGET_HOST_KEY_PATH} ${1} $(uri):${2}"
     else
         echo "sshpass -p ${TARGET_HOST_PASSWORD} ${cmd} ${1} $(uri):${2}" 
@@ -110,9 +112,9 @@ scp_to_cmd () {
 # $2 local path
 scp_from_cmd () {
     cmd="scp -r $(connect_options) "
-    if [[ -n "${BASTION_HOST}" && -n "${BASTION_HOST_USERNAME}" ]]; then
+    if [[ -n "${BASTION_HOST+x}" && -n "${BASTION_HOST_USERNAME+x}" ]]; then
         echo "${cmd} -F ssh_config target_host:${1} ${2} "
-    elif [[ -n "${TARGET_HOST_KEY_PATH}" ]]; then
+    elif [[ -n "${TARGET_HOST_KEY_PATH+x}" ]]; then
         echo "${cmd} -i ${TARGET_HOST_KEY_PATH} $(uri):${1} ${2}"
     else
         echo "sshpass -p ${TARGET_HOST_PASSWORD} ${cmd} $(uri):${1} ${2}" 
@@ -122,9 +124,9 @@ scp_from_cmd () {
 # Generate SSH command
 ssh_cmd () {
     cmd="ssh $(connect_options) "
-    if [[ -n "${BASTION_HOST}" && -n "${BASTION_HOST_USERNAME}" ]]; then
+    if [[ -n "${BASTION_HOST+x}" && -n "${BASTION_HOST_USERNAME+x}" ]]; then
         cmd+="-F ssh_config target_host "
-    elif [[ -n "${TARGET_HOST_KEY_PATH}" ]]; then
+    elif [[ -n "${TARGET_HOST_KEY_PATH+x}" ]]; then
         cmd+="-i ${TARGET_HOST_KEY_PATH} $(uri) "
     else
         cmd="sshpass -p ${TARGET_HOST_PASSWORD} ${cmd} $(uri) "
@@ -137,4 +139,29 @@ ssh_cmd () {
         cmd+=" $@"
     fi
     echo "${cmd}"
-}
\ No newline at end of file
+}
+
+# Execute command and re-try if failed
+exec_and_retry() {
+    local retries="$1"
+    local wait="$2"
+    shift
+    shift
+    local command="$@"
+
+    # Run the command, and save the exit code
+    $command
+    local exit_code=$?
+
+    # If the exit code is non-zero (i.e. command failed), and we have not
+    # reached the maximum number of retries, run the command again
+    if [[ $exit_code -ne 0 && $retries -gt 0 ]]; then
+    # Wait before retrying
+    sleep $wait
+
+    exec_and_retry $(($retries - 1)) $wait "$command"
+    else
+    # Return the exit code from the command
+    return $exit_code
+    fi
+}