open-with-local-vscode-remote

#!/bin/sh

#
# This file contains the code and instructions to set up an iTerm2 "Trigger" from a
# remote ssh session that will open up VSCode on your local machine to edit a
# file on the remote server over ssh.
#
# Author: Stuart Berg
#         https://github.com/stuarteberg
#         bergs@janelia.hhmi.org
#         https://stackoverflow.com/questions/61699447

# SETUP OVERVIEW
# --------------
# - Install the VS Code Remote Development Extension Pack
# - Ideally, setup passwordless ssh access to the remote machines you want to access
# - Place this script somewhere on your local machine (and make sure it's executable).
# - Copy the localcode() shell function below into your remote machine's .bashrc
# - Define the Trigger in iTerm2 as defined below.
#
# Notes:
#   Docs for iTerm2 Triggers: https://iterm2.com/documentation-triggers.html
#   Docs for VSCode Remote Extension: https://code.visualstudio.com/docs/remote/remote-overview
#   - CLI: https://github.com/microsoft/vscode-remote-release/issues/585#issuecomment-536580102

# iTerm2 Preferences Setup
# ------------------------
#
# In your iTerm2 preferences, set up a Trigger (Profiles > Advanced > Triggers > Edit)
#
# Regular Expression:  .*ITERM-TRIGGER-open-with-local-vscode-remote ([^ ]+) ([^ ]+) (([^ ]+ ?)+)
#             Action:  Run Command...
#         Parameters:  /path/to/this/script \1
#
# Tip: For additional feedback, try adding a duplicate entry with a "Post Notifcation" action.

# HOW TO TEST
# -----------
#
# NOTE: The new trigger will not be active for already-open terminal sessions.
#       Open a new terminal after you add the trigger to your preferences.
#
# To test it, ssh into the remote machine, and try the 'localcode' function:
#
#   localcode .
#   localcode /some/dir
#   localcode /some/file
#   localcode /some/file remote-machine-name
#
# If something is going wrong, inspect /tmp/iterm-vscode-trigger.log

#
# Put this in your remote ~/.bashrc
#

# Function to recursively find all children of a process
get_children() {
    local pid=$1
    local children=$(pgrep -P "$pid")  # Find immediate children of the given PID
    for child in $children; do
        echo -n "$child "  # Print the child PID on the same line
        get_children "$child"  # Recursively find its children
    done
}

# Main function to find and print all child PIDs of a given parent PID
find_children_pids() {
    if [ -z "$1" ]; then
        echo "Usage: find_children_pids <PID>"
        return 1
    fi

    local target_pid=$1
    # Print all child PIDs in one line
    get_children "$target_pid"
    echo  # Final newline
}

# Set this to the name of your remote machine,
# which will be serving the files you want to edit.
# (In my case, the machine is named 'submit'.)
DEFAULT_REMOTE_MACHINE_FOR_VSCODE='$(uname -n)'

localcode() {
    # Tell zsh to use bash-style arrays
    setopt ksh_arrays 2> /dev/null || true

    CMD=ITERM-TRIGGER-open-with-local-vscode-remote
    DEFAULT_MACHINE=$(eval echo ${DEFAULT_REMOTE_MACHINE_FOR_VSCODE})

    # Don't run on janelia's login nodes or submit, run on submit and don't kill sshd.
    if [[ $DEFAULT_MACHINE =~ ^e02u30.*$ || $DEFAULT_MACHINE =~ ^e02u31.*$ ]]; then
        DEFAULT_MACHINE="submit.int.janelia.org"
    fi

    MACHINE=${LOCALCODE_MACHINE-${DEFAULT_MACHINE}}

    #
    # To use this properly, your local machine should have the following ssh configuration,
    # so that VSCode can connect via port 2222 instead of the default port 22.
    #
    # Host h??u??.int.janelia.org e??u??.int.janelia.org
    #   Port 2222
    #   <and other settings here>
    #
    if [[ -z "${LOCALCODE_MACHINE}" && ! -z "${LSB_JOBID}" ]]; then
        # Launch custom ssh server on port 2222
        SSHD_PORT=$(grep -Po 'Port\s+\K[0-9]+[^#]' ~/.ssh/stuart_sshd_config)
        echo "Launching sshd on port ${SSHD_PORT}"

        /usr/sbin/sshd -f ~/.ssh/stuart_sshd_config -D >> ~/.ssh/stuart_sshd.log 2>&1 &
        SSHD_PID=$!

        echo "Setting trap to kill sshd and all its children on exit"
        trap "echo 'Killing sshd and its children: ${SSHD_PID}'; \
            children=\$(find_children_pids ${SSHD_PID}); \
            echo \${children}; \
            kill \${SSHD_PID} \${children}" EXIT
        fi

    FILENAMES=( "$@" )

    if [[ ${#FILENAMES[@]} == 0 ]]; then
        FILENAMES=.
    fi

    if [[ ${#FILENAMES[@]} == 1 && -d ${FILENAMES[0]} ]]; then
            FILENAMES[0]=$(cd ${FILENAMES[0]}; pwd)
            FTYPE=directory
    else
        # Convert filenames to abspaths
        for (( i=0; i < ${#FILENAMES[@]}; i++ )); do
            FN=${FILENAMES[i]}
            if [[ -f ${FN} ]]; then
                DIRNAME=$(cd $(dirname ${FN}); pwd)
                FILENAMES[i]=${DIRNAME}/$(basename ${FN})
                FTYPE=file
            else
                1>&2 echo "Not a valid file: ${FN}"
                exit 1
            fi
        done
    fi

    echo ${CMD} ${FTYPE} ${MACHINE} ${FILENAMES[@]}
}
export -f localcode

#
# Copy this whole file onto your local machine, or at least the following lines.
# Make sure it is executable (chmod +x /path/to/this/script)
#
trigger_vscode_remote_editing() (
    # Tell zsh to use bash-style arrays
    setopt ksh_arrays 2> /dev/null || true

    # The git extension runs 'git status -z -u' on the remote machine,
    # which takes a very long time if the remote directory is a git repo
    # with a lot of untracked files.
    # That can be fixed if you configure .gitignore appropriately,
    # but for my purposes it's easier to just disable git support when editing remote files.
    # If you want git support when using remote SSH, then comment out this line.
    # See: https://github.com/microsoft/vscode-remote-release/issues/4073
    VSCODE='/usr/local/bin/code'
    VSCODE="${VSCODE} --disable-extension vscode.git --disable-extension vscode.github --disable-extension waderyan.gitblame"
    LOGFILE=/tmp/iterm-vscode-trigger.log
    FTYPE=$1
    MACHINE=$2
    FILEPATHS=( "$@" )
    FILEPATHS=( "${FILEPATHS[@]:2}" )

    TS="["$(date "+%Y-%m-%d %H:%M:%S")"]"
    echo "${TS} Triggered: ""$@" >> ${LOGFILE}

    # https://github.com/microsoft/vscode-remote-release/issues/585#issuecomment-536580102
    if [[ "${FTYPE}" == "directory" ]]; then
        CMD="${VSCODE} --remote ssh-remote+${MACHINE} ${FILEPATHS[@]}"
        echo "${TS} ${CMD}" >> ${LOGFILE}
        ${CMD}
    elif [[ "${FTYPE}" == "file" ]]; then
        for FN in ${FILEPATHS[@]}; do
            CMD="${VSCODE} --file-uri vscode-remote://ssh-remote+${MACHINE}${FN}"
            echo "${TS} ${CMD}" >> ${LOGFILE}
            ${CMD}
        done
    else
        echo "${TS} Error: Bad arguments." >> ${LOGFILE}
        exit 1
    fi
)

export -f trigger_vscode_remote_editing

trigger_vscode_remote_editing "$@"