install-key - Generate, install, and manage SSH key rings

install-key.sh

#!/bin/bash

RED='\033[0;31m'
YELLOW='\033[1;33m'
BLUE='\033[1;96m'
NC='\033[0m' # No color

GITHUB="false"
USERNAME=""
TARGETNAME=""
PORT=""
NICKNAME=""
RING=""
KEYTYPE="ecdsa"
BITFLAG=" -b 521"
SKIP_REMOTE_INSTALL="false"

function help-text {
    echo "Usage:"
    echo "    install-key              Use the setup wizard"
    echo "    install-key -h           Display this message"
    echo "    install-key -s           Install or update the startup script and exit."
    echo "    install-key -g           GitHub mode - add a key manually or via the GitHub API"
    echo "    install-key [-x hostname] [-u username] [-p port] [-n host nickname] [-l skip remote install] [-r identity ring] [-t key type]"
    echo
    echo "Valid key types are those supported by ssh-keygen: dsa | ecdsa | ed25519 | rsa | rsa1"
    echo "Key type defaults to ecdsa with 521 bits."
    echo
}

function key-gen {

    if [ ! -d "${HOME}/.ssh" ]; then
        mkdir "${HOME}/.ssh"
    fi

    KEY_FILE="${HOME}/.ssh/${NICKNAME:-${TARGETNAME}}${RING}.key"

    echo -e "${YELLOW}-------- Generating ${KEYTYPE} key for ${USERNAME}@${TARGETNAME} --------${NC}"
    echo "It is highly recommended, albeit not required, that you use a password."
    echo "You will NOT have to type this password often. Please choose a strong one."
    echo

    ssh-keygen -t ${KEYTYPE}${BITFLAG} -f "${KEY_FILE}"
    echo
}

function add-key {

    if [ -z "${SSH_AGENT_PID}" ]; then
        echo -e "${BLUE}IMPORTANT:${NC}"
        echo -e "${NC}    Source ${HOME}/.bashrc or open a new shell to start ssh-agent and test your keys${NC}"
        echo
    else
        echo -e "${BLUE}Adding key to instance of ssh-agent...${NC}"
        echo

        $(source "${HOME}/.ssh/.install-key.env"; ssh-ring switch ${RING:1}; ssh-add "${KEY_FILE}")

        if [ $? -ne 0 ]; then
            echo
            echo "There was an issue adding ${KEY_FILE} to running instance of ssh-agent."
            echo "You may have to add it manually."
            echo "    Try: ssh-ring switch ${RING:1}; ssh-add ${KEY_FILE}"
            echo
        else
            echo
            echo "Successfully added ${KEY_FILE} to ssh-agent for identity (${RING:1})."
        fi

        echo
    fi
}

function dump-public-key {
    echo
    echo -e "${YELLOW}-------- Public Key --------${NC}"

    # Dump the public key to stdout
    cat "${KEY_FILE}.pub"
}

function add-startup-script {

    # Clean up the function from the old version
    grep -E 'start-ssh-agent\(\)' "${HOME}/.bashrc" &>/dev/null
    if [ ${?} -eq 0 ]; then

        echo "Removing old startup script..."

        cp "${HOME}/.bashrc" "${HOME}/.bashrc-install-key.bak"

        # The junk at the beginning is to read the whole file into the buffer so we can match newlines
        # Regular expression matches anything that is not a } proceeded by a newline (the closing of the function)
        sed -ri ':a;N;$!ba;s/start-ssh-agent\(\) \{([^}]|([^\n]}))*}//g' "${HOME}/.bashrc"

        if [ ${?} -ne 0 ]; then
            mv "${HOME}/.bashrc-install-key.bak" "${HOME}/.bashrc"
            echo "Cleanup of ${HOME}/.bashrc failed."
            return 1 2> /dev/null || exit 1
        fi

        # Get rid of any line referencing the old command
        sed -i ':a;N;$!ba;s/[^\n]*start-ssh-agent[^\n]*\n//g' "${HOME}/.bashrc"

        if [ ${?} -ne 0 ]; then
            mv "${HOME}/.bashrc-install-key.bak" "${HOME}/.bashrc"
            echo "Cleanup of ${HOME}/.bashrc failed."
            return 1 2> /dev/null || exit 1
        fi

        echo "Removed. Backup of .bashrc is located at ${HOME}/.bashrc-install-key.bak"
        echo
    fi

echo -e '
######## THIS FILE WILL BE OVERWRITTEN BY install-key ########

function check-permissions {

    local folder=$(realpath "${1}")
    local perms=$(stat -c "%a" "${folder}")
    local expected="${2}"

    if [ "${perms}" != "${expected}" ]; then
        echo
        echo "Permissions ${perms} on ${1} (${folder}) are too open."
        echo "Should be ${expected}. Refusing to continue."
        echo

        return 1
    fi

    local owner=$(stat -c "%U" "${folder}")

    if [ "${owner}" != "${USER}" ]; then
        echo
        echo "Owner on ${1} (${owner}) does not match user ${USER}."
        echo "Refusing to continue."
        echo

        return 1
    fi
}

current-ssh-ring() {
    SSH_DISPLAY_IDENT=${SSH_AGENT_IDENT:1} && [ -z "${SSH_DISPLAY_IDENT}" ] && SSH_DISPLAY_IDENT="default"
    echo "${SSH_DISPLAY_IDENT}"
}

list-ssh-ring() {
    for name in "${HOME}"/.ssh/.agent*; do
        # Indent if needed
        [ "${1}" == "true" ] && echo -n "  "
        sed -r "s/.env//g; s/.*.agent.?//g" <<< "${name}"
    done
}

start-ssh-ring() {
    unset SSH_AGENT_IDENT
    unset SSH_AGENT_PID
    unset SSH_AUTH_SOCK

    # Get the last used ring
    local the_path="${HOME}/.ssh/.ident.env"
    if [ -f "${the_path}" ]; then
        check-permissions "${the_path}" 600 || return 1 2> /dev/null || exit 1
        source "${the_path}"
    else
        echo "export SSH_AGENT_IDENT=${SSH_AGENT_IDENT}" > "${HOME}/.ssh/.ident.env"
        chmod 600 "${the_path}"
    fi

    local the_path="${HOME}/.ssh/.agent${SSH_AGENT_IDENT}.env"
    if [ -f "${the_path}" ]; then
        check-permissions "${the_path}" 600 || return 1 2> /dev/null || exit 1
        source "${the_path}"
    fi

    # Make sure PID is ours and is valid
    pgrep ssh-agent -u "${USER}" | grep -E "^${SSH_AGENT_PID}$" &>/dev/null
    local pid_invalid=${?}

    ssh-add -l &> /dev/null
    local agent_dead=${?}

    # PID was invalid or the socket was dead
    if [ ${pid_invalid} -ne 0 ] || [ ${agent_dead} -ne 0 ]; then

        echo "Starting ssh-agent for ring ($(current-ssh-ring))..."

        # PID is valid, but communication has been lost
        if [ ${pid_invalid} -eq 0 ]; then
            kill "${SSH_AGENT_PID}"
        fi

        rm -f "${SSH_AUTH_SOCK}"

        local the_path="${HOME}/.ssh/${SSH_AGENT_IDENT}.socket"
        if [ -f "${the_path}" ]; then
            check-permissions "${the_path}" 600 || return 1 2> /dev/null || exit 1
        fi

        # Load the keys for the selected identity
        if [ ! -z "$(ls ${HOME}/.ssh/*${SSH_AGENT_IDENT}.key 2>/dev/null)" ]; then

            # Start the agent and source its output
            eval $(ssh-agent -a "${the_path}")

            echo "
export SSH_AGENT_IDENT=${SSH_AGENT_IDENT}
export SSH_AGENT_PID=${SSH_AGENT_PID}
export SSH_AUTH_SOCK=${SSH_AUTH_SOCK}
        " > "${HOME}/.ssh/.agent${SSH_AGENT_IDENT}.env"

            chmod 600 "${HOME}/.ssh/.agent${SSH_AGENT_IDENT}.env"

            echo
            ssh-add "${HOME}"/.ssh/*"${SSH_AGENT_IDENT}".key
            echo

            local key_count=$(ssh-add -l 2> /dev/null | grep --invert-match --count "has no")
            echo "${key_count} key(s) loaded."
        else
            echo "No keys found for ($(current-ssh-ring))."
        fi
    fi
}

switch-ssh-ring() {
    export SSH_AGENT_IDENT=${1} && [ ! -z "${SSH_AGENT_IDENT}" ] && export SSH_AGENT_IDENT=".${SSH_AGENT_IDENT}"

    echo "export SSH_AGENT_IDENT=${SSH_AGENT_IDENT}" > "${HOME}/.ssh/.ident.env"
    chmod 600 "${HOME}/.ssh/.ident.env"

    start-ssh-ring
}

clean-ssh-ring() {
    if [ ! -z "$(pgrep ssh-agent -u ${USER})" ]; then
        killall -u "${USER}" ssh-agent
    fi

    rm -f "${HOME}"/.ssh/.agent*
    rm -f "${HOME}/.ssh/.ident.env"
    rm -f "${HOME}/.ssh/.socket"
    rm -f "${HOME}"/.ssh/.*.socket

    SSH_AGENT_IDENT="."
    unset SSH_AGENT_PID
    unset SSH_AUTH_SOCK
}

ssh-ring-usage() {
    echo "Commands:"
    echo "    ssh-ring"
    echo "    ssh-ring start"
    echo "    ssh-ring add|switch|sw [identity]"
    echo "    ssh-ring list"
    echo "    ssh-ring clean"
    echo
    echo "The ssh-ring command is also aliased to sshr."
}

ssh-ring() {
    if [ "${1}" == "start" ]; then
        start-ssh-ring "${@:2}"
    elif [ "${1}" == "add" ] || [ "${1}" == "switch" ] || [ "${1}" == "sw" ]; then
        switch-ssh-ring "${@:2}"
    elif [ "${1}" == "clean" ]; then
        clean-ssh-ring "${@:2}"
    elif [ "${1}" == "list" ]; then
        list-ssh-ring "${@:2}"
    elif [ "${1}" == "help" ] || [ "${1}" == "-h" ] || [ "${1}" == "--help" ]; then
        ssh-ring-usage
    else
        switch-ssh-ring "${@}"
    fi
}

alias sshr=ssh-ring

' > "${HOME}/.ssh/.install-key.env"

    chmod 600 "${HOME}/.ssh/.install-key.env"

    grep -E 'ssh-ring|sshr' "${HOME}/.bashrc" &>/dev/null
    if [ ${?} -ne 0 ]; then
        echo '
if [ -f "${HOME}/.ssh/.install-key.env" ]; then
    source "${HOME}/.ssh/.install-key.env"
    ssh-ring start
fi
' >> "${HOME}/.bashrc"

        echo "Appended auto-start script to ${HOME}/.bashrc"
    fi
}

while getopts ":hsglx:u:p:n:r:t" opt; do
    case ${opt} in
        h )
            help-text
            return 0 2> /dev/null || exit 0
            ;;
        s )
            add-startup-script
            echo "Setup script installed."
            return 0 2> /dev/null || exit 0
            ;;
        g )
            GITHUB="true"
            ;;
        x )
            TARGETNAME="${OPTARG}"
            ;;
        u )
            USERNAME="${OPTARG}"
            ;;
        p )
            PORT="${OPTARG}"
            ;;
        n )
            NICKNAME="${OPTARG}"
            ;;
        l)
            SKIP_REMOTE_INSTALL="true"
            ;;
        r )
            RING="${OPTARG}"
            ;;
        t )
            KEYTYPE="${OPTARG}"
            BITFLAG=""
            ;;
        \? )
            echo "Unknown option: -${OPTARG}" 1>&2
            help-text
            return 1 2> /dev/null || exit 1
            ;;
    esac
done

echo
echo -e "${YELLOW}-------- SSH-key generation, configuration, and installation --------${NC}"
echo

add-startup-script
source "${HOME}/.ssh/.install-key.env"

# GitHub keys do not need to be installed on a remote host
# Install using the GitHub API or print to stdout so that the user can copy/paste
if [ "${GITHUB}" == "true" ]; then

    echo -e "${YELLOW}-------- Creating SSH key for GitHub --------${NC}"
    echo

    if [ -z "${USERNAME}" ]; then
        echo "Enter GitHub username:"
        read -r USERNAME
        echo
    fi

    TARGETNAME="github.com"
    NICKNAME="${USERNAME}.github"

    key-gen

    echo
    echo "-- Enter your GitHub Personal Access Token --"
    echo "Your token will be used to authenticate to the GitHub API and add the SSH key."
    echo "If you prefer to do this yourself, leave the field blank."
    read -r -s -p "Access token: " GITHUB_TOKEN
    echo

    if [ ! -z "${GITHUB_TOKEN}" ]; then

        echo
        echo -e "${YELLOW}-------- Adding SSH key to GitHub via GitHub API --------${NC}"
        echo

        curl --fail --request POST \
        --data '{"title": "'"install-key-${HOSTNAME}"'", "key": "'"$(cat ${KEY_FILE}.pub)"'"}' \
        https://api.github.com/user/keys -K- <<< "-u ${USERNAME}:${GITHUB_TOKEN}"

        if [ ${?} -ne 0 ]; then
            echo
            echo -e "${RED}ERROR!${NC}"
            echo "There was an issue adding the key to GitHub."
            echo "Please proceed manually."
            echo
            echo "Press [Enter] to continue."

            # Let manual mode run
            GITHUB_TOKEN=""

            read
        fi

        echo
    fi

    if [ -z "${GITHUB_TOKEN}" ]; then

        dump-public-key

        echo
        echo -e "${BLUE}Paste the public key from above into the 'Key' section of:${NC}"
        echo -e "    github.com->Settings->SSH and GPG keys->New SSH Key"
        echo

    fi

    echo "Clone a new repo using: "
    echo "    git clone [email protected]:/${USERNAME}/reponame"
    echo
    echo "Switch a repo from HTTP to SSH using:"
    echo "    git remote set-url origin [email protected]:/${USERNAME}/reponame"
    echo

    # Add key to running ssh-agent or prompt user to start a new shell
    add-key

    return 0 2> /dev/null || exit 0
fi

if [ -z "${USERNAME}" ]; then
    echo "Enter username:"
    read -r USERNAME
    echo
fi

if [ -z "${TARGETNAME}" ]; then
    echo "Enter hostname:"
    read -r TARGETNAME
    echo
fi

if [ -z "${PORT}" ]; then
    echo "Enter port (blank for 22):"
    read -r PORT
    echo
fi

if [ -z "${NICKNAME}" ]; then
    echo "Enter a 'nickname' for the host (i.e. data for data.cs.purdue.edu)."
    echo "Usage after nickname will be: ssh data"
    echo "Nickname (blank will default to the hostname):"
    read -r NICKNAME
    echo
fi

if [ -z "${RING}" ]; then
    echo "Enter an identity ring (think key ring) that the key should belong to."
    echo "This allows grouping keys into separate SSH agents and helps to avoid"
    echo "authentication failures when using a large number of keys."
    echo
    echo "Existing SSH rings are:"
    ssh-ring list true
    echo
    echo "Ring (blank for default):"
    read -r RING
    echo
fi

# Append the separating dot for non-default names
if [ ! -z "${RING}" ]; then
    RING=".${RING}"
fi

# Generate the actual key
key-gen
dump-public-key

public_key_contents=$(cat "${KEY_FILE}.pub")

if [ "${SKIP_REMOTE_INSTALL}" == "true" ]; then
    echo
    echo -e "${YELLOW}Your public key has been printed above.${NC}"
    echo
    echo "Copy it and add it to ~/.ssh/authorized_keys on the remote host."
    echo "This can be done by running the following ON THE REMOTE HOST:"
    echo
    echo "    echo ${public_key_contents} >> ~/.ssh/authorized_keys"
    echo
else
    echo -e "${YELLOW}------- Installing key on remote host -------${NC}"
    echo

    ssh-copy-id -o PreferredAuthentications=keyboard-interactive,password -o PubkeyAuthentication=no -i "${KEY_FILE}" "${USERNAME}@${TARGETNAME}" -p "${PORT:-22}"

    # Make sure copy succeeded
    if [ $? -ne 0 ]; then
        echo
        echo "There was an issue installing the key to the remote host."
        echo "On the remote host, run the following command to add it manually:"
        echo
        echo "    echo ${public_key_contents} >> ~/.ssh/authorized_keys"
        echo
        SKIP_REMOTE_INSTALL=false
    fi
fi

echo

CONFIG="
Host ${NICKNAME:-${TARGETNAME}}
    Hostname ${TARGETNAME}
    IdentityFile ${KEY_FILE/#${HOME}/\~}
    Port ${PORT:-22}
    User ${USERNAME}
"
echo -e "$CONFIG" >> "${HOME}/.ssh/config"

# If newly created, will have 666
# ssh requires 600
chmod 600 "${HOME}/.ssh/config"

echo "Appended ${NICKNAME:-${TARGETNAME}} to ${HOME}/.ssh/config"

# Add key to running ssh-agent or prompt user to start a new shell
add-key

echo "Usage:"

if [ ! -z "${RING}" ]; then
    echo "    ssh-ring switch ${RING:1}"
fi

echo "    ssh ${NICKNAME:-${TARGETNAME}}"
echo

if [ "${SKIP_REMOTE_INSTALL}" == "true" ]; then
    echo -e "${YELLOW}------- Installation PARTIALLY complete -------${NC}"
    echo
    echo "You will not have SSH access until you manually add the public key to the remote host."
    echo "    Try: ssh-copy-id -i ${KEY_FILE} ${USERNAME}@${TARGETNAME}"
    echo "or see above for more detail."
else
    echo -e "${YELLOW}------- Installation complete -------${NC}"
fi