From d292ea87412fb91dd38cfdce90a2d0139c430817 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <danekhollas@gmail.com>
Date: Wed, 18 Oct 2023 16:35:16 +0100
Subject: [PATCH] Vendor load-singlesshagent.sh script (#409)

Changes of the commits are:
- vendor the `load_singlesshagent.sh` script in this repo, rather than downloading it from a super-weird readthedocs URL.
- Fixed an annoying spurious and confusing message from this script when you first enter the container and the ssh-agent is not yet running. (see third commit)
- Enabled autocompletion for aiidalab CLI app.
- Add the default private key to the `ssh-agent` after generated.
---
 stack/base/Dockerfile                         |  3 +-
 .../10_prepare-home-config.sh                 | 35 +++++-----
 stack/base/load-singlesshagent.sh             | 68 +++++++++++++++++++
 stack/lab/Dockerfile                          |  3 +
 tests/test-common.py                          |  8 +++
 5 files changed, 98 insertions(+), 19 deletions(-)
 create mode 100644 stack/base/load-singlesshagent.sh

diff --git a/stack/base/Dockerfile b/stack/base/Dockerfile
index 0867539a..52f77503 100644
--- a/stack/base/Dockerfile
+++ b/stack/base/Dockerfile
@@ -54,8 +54,7 @@ ENV AIIDA_USER_INSTITUTION Khedivial
 # Install the load-singlesshagent.sh script as described here:
 # https://aiida.readthedocs.io/projects/aiida-core/en/v2.0.0/howto/ssh.html#starting-the-ssh-agent
 # The startup of this script is configured in the before-notebook.d/20_setup-ssh.sh file.
-RUN wget --quiet --directory-prefix=/opt/bin/ \
-     "https://aiida.readthedocs.io/projects/aiida-core/en/latest/_downloads/4265ec5a42c3a3dba586dd460c0db95e/load-singlesshagent.sh"
+COPY load-singlesshagent.sh /opt/bin/
 
 # Add ~/.local/bin to PATH where the dependencies get installed via pip
 ENV PATH=${PATH}:/home/${NB_USER}/.local/bin
diff --git a/stack/base/before-notebook.d/10_prepare-home-config.sh b/stack/base/before-notebook.d/10_prepare-home-config.sh
index 5c7eea9e..a4d0ab70 100644
--- a/stack/base/before-notebook.d/10_prepare-home-config.sh
+++ b/stack/base/before-notebook.d/10_prepare-home-config.sh
@@ -6,20 +6,6 @@ if [[ ! -f  /home/${NB_USER}/.bashrc ]]; then
     cp /etc/skel/.bashrc /home/${NB_USER}/.bashrc
 fi
 
-# Set sshagent by source load-singlesshagent.sh script
-# append the command text of source to .bashrc if the script /opt/bin/load-singlesshagen.sh is present
-# and the command text is not already present in .bashrc
-header="# Load singlesshagent on shell startup."
-if [[ -f /opt/bin/load-singlesshagent.sh ]] && ! grep -q "${header}" /home/${NB_USER}/.bashrc; then
-    cat >> "/home/${NB_USER}/.bashrc" <<- EOF
-
-${header}
-if [ -f /opt/bin/load-singlesshagent.sh ]; then
-    source /opt/bin/load-singlesshagent.sh
-fi
-EOF
-fi
-
 if [[ ! -f  /home/${NB_USER}/.profile ]]; then
     cp /etc/skel/.profile /home/${NB_USER}/.profile
 fi
@@ -40,12 +26,27 @@ fi
 mkdir -p --mode=0700 /home/${NB_USER}/.ssh && \
     touch /home/${NB_USER}/.ssh/known_hosts
 
-
 if [[ ! -f  /home/${NB_USER}/.ssh/id_rsa ]]; then
   # Generate ssh key that works with `paramiko`
   # See: https://aiida.readthedocs.io/projects/aiida-core/en/latest/get_started/computers.html#remote-computer-requirements
   ssh-keygen -f /home/${NB_USER}/.ssh/id_rsa -t rsa -b 4096 -m PEM -N ''
 fi
 
-# Start the ssh-agent.
-eval `ssh-agent`
+# Set sshagent by source load-singlesshagent.sh script
+# append the command text of source to .bashrc if the script /opt/bin/load-singlesshagen.sh is present
+# and the command text is not already present in .bashrc
+header="# Load singlesshagent on shell startup."
+if [[ -f /opt/bin/load-singlesshagent.sh ]] && ! grep -q "${header}" /home/${NB_USER}/.bashrc; then
+    cat >> "/home/${NB_USER}/.bashrc" <<- EOF
+
+${header}
+if [ -f /opt/bin/load-singlesshagent.sh ]; then
+    source /opt/bin/load-singlesshagent.sh
+fi
+EOF
+fi
+
+# load the ssh-agent and add the default key generated
+# the return code can be non-zero if the ssh-agent is not running
+# which will cause the notebook to fail to start so we need to ignore the return code
+source /opt/bin/load-singlesshagent.sh || true
diff --git a/stack/base/load-singlesshagent.sh b/stack/base/load-singlesshagent.sh
new file mode 100644
index 00000000..ad7d149b
--- /dev/null
+++ b/stack/base/load-singlesshagent.sh
@@ -0,0 +1,68 @@
+# Make sure you source this script rather than executing it
+if [[ "${BASH_SOURCE[0]}" == "${0}" ]]
+then
+    echo "You need to source this script, stopping." >&2
+    exit 1
+fi
+
+# run as
+# source load-singlesshagent.sh -v
+# for verbose output
+
+load_singlesshagent() {
+    local VERBOSE
+    local SSH_ENV
+    local SSH_ADD_OUTPUT
+    local SSHADD_RETVAL
+    local NUMKEYS
+
+    VERBOSE=false
+    if [ "$1" == "-v" ]
+    then
+        VERBOSE=true
+    fi
+
+    [ "$VERBOSE" == "true" ] && echo "Single SSH agent script [verbose mode]" >&2
+    SSH_ENV="$HOME/.ssh/agent-environment"
+    # Source SSH settings, if applicable
+    if [ -r "${SSH_ENV}" ]; then
+        # don't show the output of this source command
+        source "${SSH_ENV}" 1> /dev/null
+        [ "$VERBOSE" == "true" ] && echo "- sourcing existing environment" >&2
+    else
+        [ "$VERBOSE" == "true" ] && echo "- no existing environment to source" >&2
+    fi
+
+    SSH_ADD_OUTPUT=`ssh-add -l 2> /dev/null`
+    # Needed, the later 'test' calls will replace this
+    SSHADD_RETVAL="$?"
+    # Error code: 0: there are keys; 1: there are no keys; 2: cannot contact agent
+    if [ "$SSHADD_RETVAL" == "2" ]
+    then
+        [ "$VERBOSE" == "true" ] && echo "  - unable to contact agent, creating a new one" >&2
+        (umask 066; ssh-agent > ${SSH_ENV})
+        source "${SSH_ENV}" 2> /dev/null
+    elif [ "$SSHADD_RETVAL" == "1" ]
+    then
+        [ "$VERBOSE" == "true" ] && echo "  - ssh-agent found (${SSH_AGENT_PID}), no keys (I might want to add keys here)" >&2
+        # run ssh-add to add the default generate key `id_rsa` to the agent
+        ssh-add ~/.ssh/id_rsa 2> /dev/null
+    elif [ "$SSHADD_RETVAL" == "0" ]
+    then
+        NUMKEYS=`echo "$SSH_ADD_OUTPUT" | wc -l`
+        [ "$VERBOSE" == "true" ] && echo "  - ssh-agent found (${SSH_AGENT_PID}) with $NUMKEYS keys" >&2
+    else
+        [ "$VERBOSE" == "true" ] && echo "  - ssh-add replied with return code $SSHADD_RETVAL - I don't know what to do..." >&2
+    fi
+
+    [ "$VERBOSE" == "true" ] && echo "- Debugging, listing all ssh-agents for user $NB_USER:"
+    [ "$VERBOSE" == "true" ] && ps -U "$NB_USER" | grep --color=no '[s]sh-agent'
+}
+
+# Run with the requested verbosity
+if [ "$1" == "-v" ]
+then
+    load_singlesshagent -v
+else
+    load_singlesshagent
+fi
diff --git a/stack/lab/Dockerfile b/stack/lab/Dockerfile
index 6c846483..584462c7 100644
--- a/stack/lab/Dockerfile
+++ b/stack/lab/Dockerfile
@@ -33,6 +33,9 @@ RUN mamba install --yes \
 RUN echo "aiidalab==${AIIDALAB_VERSION}" >> /opt/requirements.txt
 RUN conda config --system --add pinned_packages "aiidalab=${AIIDALAB_VERSION}"
 
+# Enable aiidalab autocompletion
+RUN echo 'eval "$(_AIIDALAB_COMPLETE=bash_source aiidalab)"' >> "${CONDA_DIR}/etc/conda/activate.d/activate_aiida_autocompletion.sh"
+
 # Install the aiidalab-home app.
 ARG AIIDALAB_HOME_VERSION
 RUN git clone https://github.com/aiidalab/aiidalab-home && \
diff --git a/tests/test-common.py b/tests/test-common.py
index a5ce02b8..acc63d8b 100644
--- a/tests/test-common.py
+++ b/tests/test-common.py
@@ -11,3 +11,11 @@ def test_verdi_status(aiidalab_exec, nb_user):
     output = aiidalab_exec("verdi status", user=nb_user).decode().strip()
     assert "Connected to RabbitMQ" in output
     assert "Daemon is running" in output
+
+
+def test_ssh_agent_is_running(aiidalab_exec, nb_user):
+    output = aiidalab_exec("ps aux | grep ssh-agent", user=nb_user).decode().strip()
+    assert "ssh-agent" in output
+
+    # also check only one ssh-agent process is running
+    assert len(output.splitlines()) == 1