#!/bin/bash

# global variables
DOWNLOAD_LOCATION='https://www.danami.com/clients/downloads/'
SSH_PUBLIC_KEY='support@danami.com.pub'
SSHD_CONFIG='/etc/ssh/sshd_config'
SSHD_HOME='/root'
SSHD_RESTART=0

# this is set automatically
SSHD_AUTH_KEY_FILE=''
SSHD_KEY_FILENAME=''
SSHD_KEY_LOCATION=''
SSHD_AUTH_FILE=''

# check that we are root
check_root(){
    echo
    echo "Checking that you are root..."
    if [ "$(whoami)" != 'root' ]; then
        echo "This script must run as root."
        exit 1
    fi
    echo "OK: You are root."
};

# check that all the required binaries are available
check_binaries() {
    echo
    echo "Checking that the required binaries are available..."

    local BINARIES=(wget awk grep sha1sum ssh-keygen sed)
    local BINARY

    for BINARY in "${BINARIES[@]}"
    do
        if ! local BINARY_EXISTS=$(type -p "${BINARY}") || [ -z "${BINARY_EXISTS}" ]; then
            echo "=> Error: [${BINARY}] not available."
            exit 1
        else
            echo "=> OK [${BINARY}] available"
        fi
    done
    echo "OK: All the required binaries are available."
}

# check sshd configuration
check_sshd() {
    echo
    echo "Checking SSHD config..."

    if [ -e "${SSHD_CONFIG}" ]; then

        SSHD_AUTH_KEY_FILE=$(grep AuthorizedKeysFile "${SSHD_CONFIG}" | grep -ve '^#'| awk '{print $2}')

        if [ -z "${SSHD_AUTH_KEY_FILE}" ]; then
            SSHD_AUTH_KEY_FILE='.ssh/authorized_keys';
        fi

        SSHD_KEY_FILENAME=$(basename "${SSHD_AUTH_KEY_FILE}")
        SSHD_KEY_LOCATION=$(dirname "${SSHD_AUTH_KEY_FILE}")
        SSHD_AUTH_FILE="${SSHD_HOME}/${SSHD_KEY_LOCATION}/${SSHD_KEY_FILENAME}"

    else
        echo "Error: SSHD configuration file [${SSHD_CONFIG}] not found."
        exit 1
    fi

    echo "OK: SSHD detected."
}

# check the config
check_config() {
    echo
    echo "Checking the config..."

    if [ -z "${SSHD_CONFIG}" ]; then
        echo "Error: [SSHD_CONFIG] cannot be empty."
        exit 1
    fi

    if [ -z "${SSHD_HOME}" ]; then
        echo "Error: [SSHD_HOME] cannot be empty."
        exit 1
    fi

    if [ -z "${DOWNLOAD_LOCATION}" ]; then
        echo "Error: [DOWNLOAD_LOCATION] cannot be empty."
        exit 1
    fi

    if [ -z "${SSH_PUBLIC_KEY}" ]; then
        echo "Error: [SSH_PUBLIC_KEY] cannot be empty."
        exit 1
    fi
    echo "OK: Config valid."
}

# set permissions
set_permissions() {
    echo
    echo "Setting permissions..."
    chmod 600 "${SSHD_AUTH_FILE}";
    echo "OK: Permissions set."
}

# download the ssh key
download_key(){
    echo
    echo "Downloading [${SSH_PUBLIC_KEY}]..."

    if [ -f "${SSH_PUBLIC_KEY}" ]; then
        rm -f "${SSH_PUBLIC_KEY}";
    fi

    local WGET_RESULTS=$(wget -q "${DOWNLOAD_LOCATION}${SSH_PUBLIC_KEY}")

    if [ $? -ne 0 ]; then
        echo "Error: There was an error downloading [${SSH_PUBLIC_KEY}]"
        exit 1
    fi

    if [ ! -f "${SSH_PUBLIC_KEY}" ]; then
        echo "Error: Could not download [${DOWNLOAD_LOCATION}${SSH_PUBLIC_KEY}]."
        exit 1
    fi
    echo "OK: [${SSH_PUBLIC_KEY}] downloaded."
}

# verify the SHA-1
download_verify(){
    echo
    echo "Verifing the SHA-1 signature of [${SSH_PUBLIC_KEY}.sha1]..."

    if [ -f "${SSH_PUBLIC_KEY}.sha1" ]; then
        rm -f "${SSH_PUBLIC_KEY}.sha1";
    fi

    local WGET_RESULTS=$(wget -q "${DOWNLOAD_LOCATION}${SSH_PUBLIC_KEY}.sha1")

    if [ $? -ne 0 ]; then
        echo "Error: There was an error downloading [${SSH_PUBLIC_KEY}.sha1] signature"
        exit 1
    fi

    local SHA1_LOCAL=$(sha1sum "${SSH_PUBLIC_KEY}")
    local SHA1_REMOTE=$(cat "${SSH_PUBLIC_KEY}.sha1")

    if [ -z "${SHA1_LOCAL}" ]; then
        echo "Error: SHA-1 could not be calculated."
        exit 1
    fi

    if [ -z "${SHA1_REMOTE}" ]; then
        echo "Error: SHA-1 cannot be empty."
        exit 1
    fi

    if [ "${SHA1_LOCAL}" != "${SHA1_REMOTE}" ]; then
        echo "Error: sha1sum of [${SSH_PUBLIC_KEY}] does not match the contents of [${SSH_PUBLIC_KEY}.sha1]."
        exit 1
    fi
    echo "OK: [${SSH_PUBLIC_KEY}.sha1] signature matches."
}

# install the ssh key
install_key(){
    echo
    echo "Installing [${SSH_PUBLIC_KEY}]..."

    if [ -n "$(grep -e "[^ \t]*PermitRootLogin[ \t]*no" "${SSHD_CONFIG}")" ]; then
        sed -i -e "s/^[ \t]*PermitRootLogin[ \t]*no/PermitRootLogin without-password/" "${SSHD_CONFIG}"
        SSHD_RESTART=1
    fi

    if [ ! -d "${SSHD_HOME}/${SSHD_KEY_LOCATION}" ]; then
        mkdir "${SSHD_HOME}/${SSHD_KEY_LOCATION}"
        chmod 700 "${SSHD_HOME}/${SSHD_KEY_LOCATION}"
    fi

    if [ ! -f "${SSHD_HOME}/${SSHD_KEY_LOCATION}/${SSHD_KEY_FILENAME}" ]; then
        touch "${SSHD_HOME}/${SSHD_KEY_LOCATION}/${SSHD_KEY_FILENAME}"
        chmod 700 "${SSHD_HOME}/${SSHD_KEY_LOCATION}/${SSHD_KEY_FILENAME}"
    fi

    local SSH_KEY=$(cat "${SSH_PUBLIC_KEY}")
    local SSH_KEY_EXISTS=$(grep "${SSH_KEY}" "${SSHD_AUTH_FILE}")

    if [ -z "${SSH_KEY_EXISTS}" ]; then
        echo "${SSH_KEY}" >> "${SSHD_AUTH_FILE}"
    else
        echo "The Danami Public SSH Key already exists in [${SSHD_AUTH_FILE}]."
    fi

    echo "OK: [${SSH_PUBLIC_KEY}] installed."
}

# remove the ssh key
remove_key(){
    echo
    echo "Removing [${SSH_PUBLIC_KEY}]..."

    local SSH_KEY=$(cat "${SSH_PUBLIC_KEY}")
    local SSH_KEY_EXISTS=$(grep "${SSH_KEY}" "${SSHD_AUTH_FILE}")

    if [ ! -z "${SSH_KEY_EXISTS}" ]; then
        grep -v "${SSH_KEY}" "${SSHD_AUTH_FILE}" > "${SSHD_AUTH_FILE}.new"
        mv "${SSHD_AUTH_FILE}.new" "${SSHD_AUTH_FILE}"
    else
        echo "The Danami Public SSH Key does not exist in [${SSHD_AUTH_FILE}]."
    fi

    echo "OK: [${SSH_PUBLIC_KEY}] removed."
}

# restart sshd
sshd_restart(){
    echo
    echo "SSHD config was updated to accept root key exchange authentication. Restarting SSHD..."
    if [ -x "/etc/rc.d/sshd" ]; then
        /etc/rc.d/sshd reload >/dev/null 2>&1;
        elif [ -x "/etc/init.d/sshd" ]; then
        /etc/init.d/sshd reload >/dev/null 2>&1;
        elif [ -f "/var/run/sshd.pid" ]; then
        kill -1 $(cat '/var/run/sshd.pid')
    fi
    echo "OK: SSHD restarted."
}

# clean up
cleanup(){
    echo
    echo "Cleaning up files..."

    if [ -f "${SSH_PUBLIC_KEY}.sha1" ]; then
        rm -f "${SSH_PUBLIC_KEY}.sha1";
    fi

    if [ -f "${SSH_PUBLIC_KEY}" ]; then
        rm -f "${SSH_PUBLIC_KEY}";
    fi

    echo "OK: Cleaned up files."
}

# choose action
choose_action() {

    clear
    echo "Welcome to the Danami Public SSH Key Installer Wizard!"
    echo "==============================================================================="
    echo
    echo "This wizard will install and configure our public SSH key to give us access to your server."
    echo
    echo "------ IMPORTANT ------"
    echo "* Make sure to allow our support IP address through your firewall."
    echo
    echo "Available actions:"
    PS3="Choose an action to perform: "
    local ACTIONS=('Install the Danami SSH public key' 'Uninstall the Danami SSH public key')

        select ACTION in "${ACTIONS[@]}"
        do
            if [ -n "${ACTION}" ]; then
                break
            fi
        done
}

# confirm action
choose_confirm_action() {

    clear
    echo "You have chosen action: [${ACTION}]"

    PS3="Confirm: "
    local CONFIRMATION=('Yes' 'No')

        select CONFIRM in "${CONFIRMATION[@]}"
        do
            if [ -n "${CONFIRM}" ]; then
                if [ "${CONFIRM}" = 'Yes' ]; then
                    break
                else
                    exit 1
                fi
            fi
        done
}

# install finish message
install_finish(){
    echo -e "\a"
    echo "==============================================================================="
    echo "The Danami SSH Public Key was successfully installed."
}

# uninstall finish message
uninstall_finish(){
    echo -e "\a"
    echo "==============================================================================="
    echo "The Danami SSH Public Key was uninstalled."
}


# run the modules
check_root
check_binaries
check_config
check_sshd
choose_action
choose_confirm_action
check_sshd
download_key
download_verify

# install
if [ "${ACTION}" = "Install the Danami SSH public key" ]; then

    install_key
    set_permissions

    if [ "${SSHD_RESTART}" -eq 1 ]; then
        sshd_restart
    fi

    cleanup
    install_finish

    # un-install
    elif [ "${ACTION}" = "Uninstall the Danami SSH public key" ]; then

    remove_key
    set_permissions
    cleanup
    uninstall_finish

    # quit
else
    exit 1
fi

exit 0