#!/usr/bin/env bash

# Prepare a LXC container for remote management by adding authorized SSH keys
# to its 'root' account. The script will ensure that OpenSSH is installed in
# the container.

set -o nounset -o pipefail -o errexit

readonly SCRIPT="$(basename "${0}")"
readonly CONTAINER="${1:-}"

if [ -n "${SUDO_USER:-}" ] ; then
    readonly SUDO_USER_HOME="$(getent passwd "${SUDO_USER}" | cut -d: -f6)"
    readonly SUDO_AUTHORIZED_KEYS="${SUDO_USER_HOME}/.ssh/authorized_keys"

    # Support for SSH keys from LDAP directory
    if [ -x "/etc/ssh/authorized_keys_lookup" ] ; then
        readonly LDAP_AUTHORIZED_KEYS="$(/etc/ssh/authorized_keys_lookup "${SUDO_USER}")"
    fi
else
    readonly SUDO_AUTHORIZED_KEYS=""
fi

container_exists () {
    local container
    container="${1}"

    if lxc-ls -1 | grep -q -E "^${container}$" ; then
        return 0
    else
        return 1
    fi
}

container_is_running () {
    local container
    container="${1}"

    if [ "$(lxc-info -n "${container}" | grep -E "^State:\\s" | awk '{print $2}')" = "RUNNING" ] ; then
        return 0
    else
        return 1
    fi
}

prepare_openssh () {
    local container
    container="${1}"

    if ! lxc-attach -n "${container}" -- test -x /usr/sbin/sshd ; then
        if lxc-attach -n "${container}" -- test -f /etc/debian_version ; then
            printf "Installing OpenSSH service...\\n"
            lxc-attach -n "${container}" -- apt-get update
            lxc-attach -n "${container}" -- apt-get -y --no-install-recommends install openssh-server
        fi
    fi
}

prepare_authorized_keys () {
    local container
    container="${1}"
    local authorized_keys
    authorized_keys="${2}"

    local -a authorized_keys_content
    if [ -f "${authorized_keys}" ] ; then
        readarray -t authorized_keys_content <<< "$(< "${authorized_keys}")"
    else
        readarray -t authorized_keys_content <<< "${authorized_keys}"
    fi

    if ! lxc-attach -n "${container}" -- test -d /root/.ssh ; then
        lxc-attach -n "${container}" -- mkdir -p -m 0700 /root/.ssh
        lxc-attach -n "${container}" -- touch /root/.ssh/authorized_keys
        lxc-attach -n "${container}" -- chmod 0600 /root/.ssh/authorized_keys
    fi
    if [ -f "${authorized_keys}" ] ; then
        printf "Adding authorized SSH keys from '%s'...\\n" "${authorized_keys}"
    else
        printf "Adding authorized SSH key...\\n"
    fi
    for key_line in "${authorized_keys_content[@]}" ; do
        if ssh-keygen -l -f - <<< "${key_line}" > /dev/null ; then
            if ! lxc-attach -n "${container}" -- grep -Fxq "${key_line}" /root/.ssh/authorized_keys ; then
                ssh-keygen -l -f - <<< "${key_line}"
                lxc-attach -n "${container}" -- tee -a /root/.ssh/authorized_keys <<< "${key_line}" 1>/dev/null
            fi
        fi
    done
}

print_usage () {
    cat <<EOF
${SCRIPT}: add authorized SSH keys to an LXC container

Usage: ${SCRIPT} <container-name>

The keys will be taken from files:
    \${SUDO_USER}/.ssh/authorized_keys
If enabled, LDAP directory will be checked for SSH authorized_keys
EOF
}

error_msg () {
    printf "%s\\n" "${*}" >&2
}

main () {
    local container
    container="${CONTAINER}"
    local -a authorized_keys
    authorized_keys=( "${SUDO_AUTHORIZED_KEYS}" )
    local ldap_authorized_keys
    ldap_authorized_keys="${LDAP_AUTHORIZED_KEYS:-}"

    if [ -n "${container}" ] ; then
        if container_exists "${container}" ; then
            if ! container_is_running "${container}" ; then
                printf "Container '%s' is not running, starting...\\n" "${container}"

                if pidof systemd > /dev/null 2>&1 ; then
                    systemctl enable "lxc@${container}.service"
                    systemctl start "lxc@${container}.service"
                else
                    lxc-start -n "${container}" -d
                fi
                lxc-wait -n "${container}" -s RUNNING
            fi
            prepare_openssh "${container}"
            for key_file in "${authorized_keys[@]}" ; do
                if [ -f "${key_file}" ] ; then
                    prepare_authorized_keys "${container}" "${key_file}"
                fi
            done
            if [ -n "${ldap_authorized_keys}" ] ; then
                prepare_authorized_keys "${container}" "${ldap_authorized_keys}"
            fi
        else
            error_msg "Error: container '${container}' doesn't exist"
            return 1
        fi
    else
        print_usage
        return 1
    fi
}

main
