#!/bin/bash
#
# Copyright 2019 Deutsches Klimarechenzentrum GmbH
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above
#    copyright notice, this list of conditions and the following
#    disclaimer in the documentation and/or other materials provided
#    with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.
#
#
# start-vnc
#
# This script is intended to be used on your local workstation running
# Ubuntu or Fedora Linux or macOS (tested). Other Unix flavors may
# work as well. It allows you to connect to one of DKRZ's
# visualization nodes to work remotely on your visualizations.
#
# Technically, it starts a VNC server on one of the GPU nodes by
# submitting a generated job script. A local vncviewer client is used
# to connect to the server over an encypted ssh tunnel. After the
# client shuts down, the server job is terminated. TigerVNC, TightVNC,
# and TurboVNC are known to work.
#
# In case of problems contact Mathis Rosenhauer <rosenhauer@dkrz.de>.
#

set -eufo pipefail

# Default settings
#
# You can change the settings here or override them with command line
# options. SVNC_ACCTCODE is the only parameter you have to provide if
# your local username matches your username on the frontend and if
# vncviewer is installed in the search path (tight/tigervnc on Linux
# for example).

# Project account code
SVNC_ACCTCODE=""

# LDAP username
#
# Specify your username on the login node if it is not your local
# username.
SVNC_USERNAME="$(id -un)"

# Path to local vncviewer
#
# If your local vnc client is not in the search path or named
# differently (e.g. TurboVNC, macOS), then change this parameter.
SVNC_CLIENT="vncviewer"

# Server options
#
# More options for the vncserver. TurboVNC on A Mac may produce
# a 'javax.net.ssl.SSLHandshakeException' in this case adding
# ' -securitytypes none' to SVNC_SERVER_OPTIONS may help.
SVNC_SERVER_OPTIONS="-geometry 1920x1200"

# Session run time in minutes
SVNC_RUNTIME=240

# Number of GPU nodes reserved. Default is the special value "half"
# which allocates 24 cores of a node. To allocate one or more nodes,
# set to the integer number of nodes.
SVNC_NODES="half"

# Quality of service
SVNC_QOS="normal"

# Constraint for feature selection (GPU and/or memory)
SVNC_CONSTRAINT=""

# Uncomment this if you use TurboVNC as vncviewer.
#
# TurboVNC will work without this option but then it will use its own
# implementation of ssh which dosen't support public key auth among other
# things.
#
# readonly SVNC_CLIENT_OPTIONS="-extssh"

# Frontend host
#
# Must be directly accessible from client. The frontend and the node
# where vncserver is running need a shared home file system. You will
# have to change this for other sites than DKRZ (along with many other
# things).
readonly SVNC_FRONTEND_HOST="mistral.dkrz.de"

# Copy vncpassword temporarily to the local workstation
#
# If your vncviewer client has trouble with TurboVNC's password format
# then disable this option (set to "false" or comment).
readonly SVNC_PASSWORD=true

# Kill vncserver job running if script ends for some reason
readonly SVNC_KILLSERVER=true


clean_up () {
    trap - ERR EXIT
    set +e

    if [[ "${SVNC_KILLSERVER}" = true ]]; then
        if [[ -n ${vnc_host:-} ]]; then
            echo "Killing vncserver :${vnc_display} on ${vnc_host}."
            ssh_frontend "ssh ${vnc_host} \"/opt/TurboVNC/bin/vncserver \
                         -kill :${vnc_display}\""
        fi

        if [[ -n ${job_id:-} ]]; then
            echo "Removing job ${job_id}."
            ssh_frontend "scancel -Q ${job_id}; sleep 1; \
                         rm -f .startvnc/out.${job_id}"
        else
            echo "Job ID not available. Make sure the vncjob is not running!"
            ssh_frontend "squeue -u ${SVNC_USERNAME}"
        fi
    else
        echo "Kept job ${job_id} running."
    fi

    ssh_frontend "" "-O exit"
    rmdir "${ssh_socket_dir}"

    # Remove local vnc PasswordFile
    if [[ "${SVNC_PASSWORD}" = true ]]; then
        rm -f "vnc_passwd"
    fi
    exit
}

usage () {
    cat <<EOF
Usage: $(basename "$0") [OPTION]

Available values for OPTION include:

  -A acctcode  acctcode of job
  -C list      constraint for feature selection
  -n nodes     number of nodes
  -q qos       job qos
  -s options   addidtional options for vncserver
  -t time      job runtime
  -u username  use username for login
  -v path      path to vncviewer program
  -x options   addidtional options for vncviewer

EOF
}

ssh_frontend () {
    local command="$1"
    local extra_options="${2:-}"
    local options
    options="${extra_options} -o ForwardX11=no \
            -o ControlPath=${ssh_socket_dir}/control:%h:%p:%r"
    ssh ${options} "${SVNC_USERNAME}@${SVNC_FRONTEND_HOST}" "${command}"
}

parse_options () {
    local option
    while getopts 'A:C:n:q:s:t:u:v:x:' option; do
        case ${option} in
            A) SVNC_ACCTCODE="$OPTARG"
               ;;
            C) SVNC_CONSTRAINT="$OPTARG"
               ;;
            n) SVNC_NODES="$OPTARG"
               ;;
            q) SVNC_QOS="$OPTARG"
               ;;
            s) SVNC_SERVER_OPTIONS="$OPTARG"
               ;;
            t) SVNC_RUNTIME="$OPTARG"
               ;;
            u) SVNC_USERNAME="$OPTARG"
               ;;
            v) SVNC_CLIENT="$OPTARG"
               ;;
            x) SVNC_CLIENT_OPTIONS="$OPTARG"
               ;;
            ?) usage
            exit 1
            ;;
        esac
    done

    readonly SVNC_ACCTCODE
    readonly SVNC_CONSTRAINT
    readonly SVNC_SERVER_OPTIONS
    readonly SVNC_QOS
    readonly SVNC_RUNTIME
    readonly SVNC_USERNAME
    readonly SVNC_CLIENT
    readonly SVNC_CLIENT_OPTIONS

    if [[ -z ${SVNC_ACCTCODE} ]]; then
        printf "ERROR: Please specify an acctcode.\n\n" >&2
        usage
        exit 1
    fi
}

prepare_vncserver () {
    ssh_frontend "mkdir -p .startvnc"
    if ! ssh_frontend "test -s .vnc/passwd"; then
        echo "No VNC password found. Please set now."
        echo "Do not use your LDAP password. Eight characters maximum."
        ssh_frontend "salloc -Q -n1 -pgpu -A${SVNC_ACCTCODE} -- /bin/bash -c \
                     'ssh -tt \$SLURM_JOB_NODELIST -- mkdir -p .vnc \
                     && /opt/TurboVNC/bin/vncpasswd'" "-t"
    fi
}

submit_vnc_job () {
    local sbatch_resources
    if [[ ${SVNC_NODES} = "half" ]]; then
        sbatch_resources="#SBATCH --ntasks=24"
    else
        sbatch_resources="#SBATCH --nodes=${SVNC_NODES}"
        sbatch_resources+=$'\n#SBATCH --exclusive'
    fi
    ssh_frontend "cd .startvnc && sbatch" <<EOF
#!/bin/bash -l
#SBATCH --job-name=vncserver
#SBATCH --partition=gpu
#SBATCH --constraint=${SVNC_CONSTRAINT}
${sbatch_resources}
#SBATCH --qos=${SVNC_QOS}
#SBATCH --time=${SVNC_RUNTIME}
#SBATCH --account=${SVNC_ACCTCODE}
#SBATCH --output=out.%j
#SBATCH --parsable
#SBATCH --dkrzepilog=0
cd \${HOME}
/opt/TurboVNC/bin/vncserver -fg -localhost ${SVNC_SERVER_OPTIONS}
EOF
}

get_vnc_host_and_display () {
    local job_id="$1"
    ssh_frontend "/bin/bash -s" <<EOF
logfile=\${HOME}/.startvnc/out.${job_id}
while [[ -z \${host_and_display} ]]; do
    sleep 1
    if [[ -f \${logfile} ]]; then
        host_and_display="\$(grep -Po "started on display \Kmg[0-9]+:[0-9]+" \
                          \${logfile})"
    fi
    printf "." >&2
done
printf "\n" >&2
echo \${host_and_display}
EOF
}

main () {
    parse_options "$@"
    trap clean_up INT QUIT TERM ERR EXIT

    mkdir -p "${HOME}/.ssh"
    ssh_socket_dir="$(mktemp -d "${HOME}/.ssh/socket.XXXXX")"
    ssh_frontend "" "-MNf"
    prepare_vncserver

    echo "Submitting vncserver job."
    job_id="$(submit_vnc_job)"

    printf "Waiting for job ${job_id} to start" >&2
    local host_and_display
    host_and_display="$(get_vnc_host_and_display "$job_id")"
    vnc_host=${host_and_display%:*}
    vnc_display=${host_and_display#*:}
    echo "Vncserver started on node ${vnc_host}.dkrz.de display \
:${vnc_display}."

    # WSL doesn't seem to detect used ports so randomize
    # local port to reduce risk of masking.
    local vnc_port_local=$((5900 + RANDOM % 100))
    local vnc_port_remote=$(( 5900 + vnc_display ))
    until ssh -o ForwardX11=no \
              -o StrictHostKeyChecking=accept-new \
              -L "${vnc_port_local}:localhost:${vnc_port_remote}" \
              -Nf \
              "${SVNC_USERNAME}@${vnc_host}.dkrz.de"; do
        vnc_port_local=$((5900 + RANDOM % 100))
        echo "Trying local port ${vnc_port_local}."
    done
    local client_options
    client_options="${SVNC_CLIENT_OPTIONS:-}"
    if [[ "${SVNC_PASSWORD}" = true ]]; then
        echo "Fetching password from frontend."
        ssh_frontend "cat .vnc/passwd" > vnc_passwd
        client_options+=" -passwd vnc_passwd"
    fi
    echo "Connecting vncviewer to ${vnc_host}.dkrz.de"
    "${SVNC_CLIENT}" ${client_options} :$(( vnc_port_local - 5900 ))
}

main "$@"