From b03718cf7dea8ec188be84932aee4f13aa048c98 Mon Sep 17 00:00:00 2001
From: Zachary Lentz <zlentz@slac.stanford.edu>
Date: Wed, 27 Mar 2024 10:31:03 -0700
Subject: [PATCH] ENH/MAINT: rework the error handling to avoid bad states

---
 scripts/bootstrap_plc.sh     | 5 ++---
 scripts/dry_run.sh           | 2 ++
 scripts/provision_plc.sh     | 7 ++++---
 scripts/setup_new_plc.sh     | 5 ++---
 scripts/ssh_agent_helper.sh  | 3 +++
 scripts/update_admin_pass.sh | 5 ++---
 6 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/scripts/bootstrap_plc.sh b/scripts/bootstrap_plc.sh
index fcc2afb..289fa81 100755
--- a/scripts/bootstrap_plc.sh
+++ b/scripts/bootstrap_plc.sh
@@ -50,6 +50,8 @@ fi
 
 # Register the ssh key with the ssh agent if needed
 source "${THIS_DIR}/ssh_agent_helper.sh"
+# Stop the ssh agent at exit if we started it here
+trap ssh_agent_helper_cleanup EXIT
 
 # Send the public key to the plc, if it has not already been done
 ssh-copy-id -i "${SSH_KEY_FILENAME}" -o PreferredAuthentications=keyboard-interactive "${USERNAME}@${HOSTNAME}"
@@ -91,6 +93,3 @@ scp -F "${SSH_CONFIG}" -i "${SSH_KEY_FILENAME}" -r "${SOURCE_DIR}" "${USERNAME}@
 
 # Run the local install version of the bootstrap playbook
 ansible-playbook "${ANSIBLE_ROOT}/tcbsd-bootstrap-from-local-playbook.yaml" --extra-vars "target=${HOSTNAME} ansible_ssh_private_key_file=${SSH_KEY_FILENAME}" --ask-become-pass "$@"
-
-# Stop the ssh agent if we started it here
-ssh_agent_helper_cleanup
diff --git a/scripts/dry_run.sh b/scripts/dry_run.sh
index 56dab27..d5257ed 100755
--- a/scripts/dry_run.sh
+++ b/scripts/dry_run.sh
@@ -10,6 +10,8 @@
 #   $ ./dry_run.sh tst_all
 #
 # Groups are defined in the inventory file.
+set -e
+
 if [ -z "${1}" ]; then
   echo "Ansible target required"
   exit 1
diff --git a/scripts/provision_plc.sh b/scripts/provision_plc.sh
index 6834d1d..2590515 100755
--- a/scripts/provision_plc.sh
+++ b/scripts/provision_plc.sh
@@ -10,6 +10,8 @@
 #   $ ./provision_plc.sh tst_all
 #
 # Groups are defined in the inventory file.
+set -e
+
 if [ -z "${1}" ]; then
   echo "Ansible target required"
   exit 1
@@ -27,12 +29,11 @@ source "${THIS_DIR}"/activate_python.sh
 
 # Register the ssh key with the ssh agent if needed
 source "${THIS_DIR}/ssh_agent_helper.sh"
+# Stop the ssh agent at exit if we started it here
+trap ssh_agent_helper_cleanup EXIT
 
 # Run the provision playbook
 ansible-playbook "${ANSIBLE_ROOT}/tcbsd-provision-playbook.yaml" --extra-vars "target=${TARGET} ansible_ssh_private_key_file=${SSH_KEY_FILENAME}" --ask-become-pass "$@"
 
-# Stop the ssh agent if we started it here
-ssh_agent_helper_cleanup
-
 # Prompt to update deployment docs
 "${THIS_DIR}"/docs_prompt.sh
diff --git a/scripts/setup_new_plc.sh b/scripts/setup_new_plc.sh
index d8774f1..74e007f 100755
--- a/scripts/setup_new_plc.sh
+++ b/scripts/setup_new_plc.sh
@@ -18,10 +18,9 @@ THIS_DIR="$(dirname "${THIS_SCRIPT}")"
 
 # Register the ssh key with the ssh agent if needed
 source "${THIS_DIR}/ssh_agent_helper.sh"
+# Stop the ssh agent at exit if we started it here
+trap ssh_agent_helper_cleanup EXIT
 
 # Run both playbooks and one-time pre-playbook setup
 "${THIS_DIR}"/bootstrap_plc.sh "${1}"
 "${THIS_DIR}"/provision_plc.sh "${1}"
-
-# Stop the ssh agent if we started it here
-ssh_agent_helper_cleanup
diff --git a/scripts/ssh_agent_helper.sh b/scripts/ssh_agent_helper.sh
index 8c4ff29..f904fa9 100644
--- a/scripts/ssh_agent_helper.sh
+++ b/scripts/ssh_agent_helper.sh
@@ -7,6 +7,7 @@
 # Expected usage:
 #
 # source ssh_agent_helper.sh
+set -e
 
 SSH_KEY_FILENAME="${HOME}/.ssh/tcbsd_key_rsa"
 export SSH_KEY_FILENAME
@@ -28,6 +29,8 @@ ssh_agent_helper_cleanup() {
     fi
 }
 export ssh_agent_helper_cleanup
+# Clean up immediately if something in this script fails
+trap ssh_agent_helper_cleanup ERR
 
 # Create an ssh key, if it does not already exist
 if [ ! -f "${SSH_KEY_FILENAME}" ]; then
diff --git a/scripts/update_admin_pass.sh b/scripts/update_admin_pass.sh
index 1b2eb63..356b881 100755
--- a/scripts/update_admin_pass.sh
+++ b/scripts/update_admin_pass.sh
@@ -26,11 +26,10 @@ SSH_CONFIG="${ANSIBLE_ROOT}/ssh_config"
 
 # Register the ssh key with the ssh agent if needed
 source "${THIS_DIR}/ssh_agent_helper.sh"
+# Stop the ssh agent at exit if we started it here
+trap ssh_agent_helper_cleanup EXIT
 
 for HOSTNAME in "$@"; do
     echo "Logging into ${HOSTNAME}"
     ssh -F "${SSH_CONFIG}" -i "${SSH_KEY_FILENAME}" -t "${USERNAME}@${HOSTNAME}" passwd
 done
-
-# Stop the ssh agent if we started it here
-ssh_agent_helper_cleanup