Skip to content

Commit

Permalink
Vendor load-singlesshagent.sh script (#409)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
danielhollas authored Oct 18, 2023
1 parent 2d35cca commit d292ea8
Show file tree
Hide file tree
Showing 5 changed files with 98 additions and 19 deletions.
3 changes: 1 addition & 2 deletions stack/base/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
35 changes: 18 additions & 17 deletions stack/base/before-notebook.d/10_prepare-home-config.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
68 changes: 68 additions & 0 deletions stack/base/load-singlesshagent.sh
Original file line number Diff line number Diff line change
@@ -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
3 changes: 3 additions & 0 deletions stack/lab/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -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 && \
Expand Down
8 changes: 8 additions & 0 deletions tests/test-common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit d292ea8

Please sign in to comment.