#!/bin/sh
set -e
set -u

main() {
    cmd_sudo=''
    if command -v sudo > /dev/null; then
        my_id="$(id -u)"
        if test "0" != "${my_id}"; then
            cmd_sudo='sudo'
        fi
    fi

    cmd_sed="${cmd_sudo} sed -i -E"
    my_bsd_sed=''
    if ! sed -V 2>&1 | grep -q 'GNU'; then
        cmd_sed="${cmd_sudo} sed -i .prohibit-password-bak -E"
        my_bsd_sed='true'
    fi

    my_changes=''

    # PasswordAuthentication
    echo ""
    echo "Checking for 'PasswordAuthentication no'..."
    echo "    grep '^\s*PasswordAuthentication\s*no' /etc/ssh/sshd_config"
    my_allow_passwords="$(
        grep -q \
            '^\s*PasswordAuthentication\s*no\s*$' \
            /etc/ssh/sshd_config ||
            echo 'true'
    )"
    if test -n "${my_allow_passwords}"; then
        $cmd_sed \
            's/#*[[:space:]]*PasswordAuthentication[[:space:]]*(yes|no)[[:space:]]*$/PasswordAuthentication no/g' \
            /etc/ssh/sshd_config

        if ! grep -q '^PasswordAuthentication no$' /etc/ssh/sshd_config; then
            echo ""
            echo "ERROR"
            echo "    failed to update '/etc/ssh/sshd_config'"
            echo ""
            return 1
        fi
        echo "    RESTART REQUIRED: disabled user password login"
        my_changes='true'
    else
        echo "    PASS: passwords are NOT allowed"
    fi

    # PermitRootLogin
    echo ""
    echo "Checking for 'PermitRootLogin prohibit-password'..."
    echo "    grep -E '^\s*PermitRootLogin\s*(no|prohibit-password)' /etc/ssh/sshd_config"
    my_allow_root="$(
        grep -q -E \
            '^\s*PermitRootLogin\s*(no|prohibit-password|without-password)\s*$' \
            /etc/ssh/sshd_config ||
            echo 'true'
    )"
    if test -n "${my_allow_root}"; then
        $cmd_sed \
            's/#*[[:space:]]*PermitRootLogin[[:space:]]*(yes|no|prohibit-password|without-password)[[:space:]]*$/PermitRootLogin prohibit-password/g' \
            /etc/ssh/sshd_config
        if ! grep -q -E \
            '^PermitRootLogin prohibit-password$' \
            /etc/ssh/sshd_config; then
            echo ""
            echo "ERROR"
            echo "    failed to update '/etc/ssh/sshd_config'"
            echo ""
            return 1
        fi
        echo "    RESTART REQUIRED: disabled root password login"
        my_changes='true'
    else
        echo "    PASS: 'root' can NOT login via password"
    fi

    # UsePAM
    echo ""
    echo "Checking for 'UsePAM no' (macOS only)..."
    # shellcheck disable=SC2016
    echo '    test "$(uname -s)" = "Darwin" &&'
    echo "    grep '^\s*UsePAM\s*no' /etc/ssh/sshd_config"
    if test "$(uname -s)" = "Darwin"; then
        my_allow_pam="$(
            grep -q \
                '^\s*UsePAM\s*no' \
                /etc/ssh/sshd_config ||
                echo 'true'
        )"
        if test -n "${my_allow_pam}"; then
            $cmd_sed \
                's/#*[[:space:]]*UsePAM[[:space:]]*(yes|no)[[:space:]]*$/UsePAM no/g' \
                /etc/ssh/sshd_config
            if ! grep -q '^UsePAM no$' /etc/ssh/sshd_config; then
                echo ""
                echo "ERROR"
                echo "    failed to update '/etc/ssh/sshd_config'"
                echo ""
                return 1
            fi
            echo "    RESTART REQUIRED: disabled macOS password login"
            my_changes='true'
        else
            echo "    PASS: passwords are NOT allowed"
        fi
    fi

    if test -z "${my_changes}"; then
        echo ""
        echo ""
        echo "All checks pass. No changes necessary."
        echo ""
        return 0
    fi

    if test -n "${my_bsd_sed}"; then
        $cmd_sudo rm -f /etc/ssh/sshd_config.prohibit-password-bak
    fi

    echo ""
    echo "##################################################################"
    echo "#                                                                #"
    echo "#  IMPORTANT                                                     #"
    echo "#                                                                #"
    echo "#  READ CAREFULLY                                                #"
    echo "#                                                                #"
    echo "#  (or get locked out)                                           #"
    echo "#                                                                #"
    echo "##################################################################"
    echo ""

    echo ""
    echo "1. TEST SSH KEYS"
    echo ""
    echo "    Be sure that you can login as an admin user with ssh keys"
    echo ""
    echo "2. RESTART SSH with one of the following:"
    echo ""
    echo "    systemctl restart sshd   # Ubuntu / Debian"
    echo "    rc-service sshd restart  # Alpine / Gentoo"
    echo "    killall sshd             # all others"
    echo ""
}

main "${1:-app}" "${2:-}"
