#!/bin/bash

declare -a dips
declare -a sshnode
declare -a expnode

duser="root"
dpasswd=""

ssh_option1="-o StrictHostKeyChecking=no"
ssh_option2="-o PreferredAuthentications=publickey"
ssh="/usr/bin/ssh $ssh_option1"
scp="/usr/bin/scp $ssh_option1 -q"

keygen_cmd=""
private_keygen_file=""
pub_keygen_file=""
authorized_file=""
pub_keygen=""
pub_rsa="/tmp/id_rsa.pub"

detail=false

# Crate id_rsa, id_rsa_pub, if not exists id_rsa
function keygen_check_ssh
{
    local node="$1"
    if ! ($ssh $node "/usr/bin/test -f $private_keygen_file" 2> /dev/null); then
        if ! ($ssh $node "$keygen_cmd"); then
            echo "[$node]: Failed to create public key with ssh"
            return 1
        fi
    fi

    if $detail; then
        echo "    $1 check Done"
    fi
    return $?
}

# From node get id_rsa.pub save /tmp
function keygen_copy_ssh {
    local node="$1"
    $scp $node:$pub_keygen_file /tmp/ 2> /dev/null
    if [ $? -ne 0 ]; then
        echo "[$node]: Failed to get public key"
        return 1
    fi

    send_keygen "$node"
    return $?
}

# Expect run shell command
function expect_run
{
    local cmd=$1
    dpasswd=$(echo $dpasswd | sed 's/;/\\;/')
    /usr/bin/expect -c "set timeout -1;
        spawn $cmd;
        expect {
            *password:* {send -- $dpasswd\r;
                expect {
                    *denied* {exit 1;}
                    eof
                }
            }
            eof {exit 0;}
        }";
    return $?
}

# Crate id_rsa, id_rsa_pub, if not exists
function keygen_check_exp {
    local node="$1"
    local errmsg=""
    local keygen_file_cmd="/usr/bin/test -f $private_keygen_file || $keygen_cmd"
    errmsg=$(expect_run "$ssh $node \"$keygen_file_cmd\"" | grep -Ev '^spawn.*$|^.*password.*$|^Authorized.*$|^$'|tr -d '\r')
    if [[ "x$errmsg" != "x" ]]; then
        echo "[$node]: Failed to create public key with expect"
        return 1
    fi

    if $detail; then
        echo "    $1 check Done"
    fi
    return $?
}

# From node get id_rsa.pub save /tmp
function keygen_copy_exp
{
    local node="$1"
    local errmsg=""
    # Send id_rsa.pub to other host
    errmsg=$(expect_run "$scp $node:$pub_keygen_file /tmp/" | grep -Ev '^spawn.*$|^.*password.*$|^Authorized.*$|^$'|tr -d '\r')
    if [[ "x$errmsg" != "x" ]]; then
        echo "[$node]: Failed to get public key"
        return 1
    fi

    send_keygen "$node"
    return $?
}

# Send $node id_rsa.pub to all host .ssh/authorized_keys
function send_keygen
{
    local node="$1"
    local pub_keygen="$(/bin/cat $pub_rsa)"
    local send_cmd="grep -qw '$pub_keygen' $authorized_file 2> /dev/null || /bin/echo ${pub_keygen} >> $authorized_file"
    local errmsg=""

    # shellcheck disable=SC2068
    for host in ${sshnode[@]}; do
        if ! ($ssh $host "${send_cmd}" &> /dev/null); then
            echo "[$node]: Failed to send public key to $host"
            return 1
        fi
    done
    # shellcheck disable=SC2068
    for host in ${expnode[@]}; do
        errmsg=$(expect_run "$ssh $host \"$send_cmd\"" | grep -Ev '^spawn.*$|^.*password.*$|^.*Warning:.*(RSA).*$|^$|^Authorized.*$'|tr -d '\r')
        if [[ "x$errmsg" != "x" ]]; then
            echo "$errmsg"
            echo "[$node]: Failed to send public key to $host"
            return 1
        fi
    done

    if $detail; then
        echo "    $1 keygen Done."
    fi
    return 0
}

# Set up mutual trust
function setup_mutual_trust
{
    if [[ "$duser" == "root" ]]; then
        private_keygen_file="/root/.ssh/id_rsa"
        pub_keygen_file="/root/.ssh/id_rsa.pub"
        authorized_file="/root/.ssh/authorized_keys"
        keygen_cmd="/usr/bin/ssh-keygen -q -N '' -t rsa -f /root/.ssh/id_rsa"
    else
        private_keygen_file="/home/$duser/.ssh/id_rsa"
        pub_keygen_file="/home/$duser/.ssh/id_rsa.pub"
        authorized_file="/home/$duser/.ssh/authorized_keys"
        keygen_cmd="/usr/bin/ssh-keygen -q -N '' -t rsa -f /home/$duser/.ssh/id_rsa"
    fi

    # check all node private_keygen_file
    if $detail; then
        echo -e "\nCheck Start"
    fi

    # shellcheck disable=SC2068
    for host in ${dips[@]}; do
        if ($ssh $ssh_option2 $host "/bin/date" &> /dev/null); then
            echo "check private key files of $host by ssh"
            sshnode+=("$host")
            keygen_check_ssh $host
            if [ $? -ne 0 ]; then
                echo "Mutual trust failed"
                return 1
            fi
        else
            echo "check private key of $host by expect"
            expnode+=("$host")
            keygen_check_exp $host
            if [ $? -ne 0 ]; then
                echo "Mutual trust failed"
                return 1
            fi
        fi
    done

    if $detail; then
        echo "Check Done."
        echo -e "\nKeygen Start"
    fi

    # get pub_keygen_file from ssh node
    # shellcheck disable=SC2068
    for host in ${sshnode[@]}; do
        echo "get pub keygen file from ssh $host"
        keygen_copy_ssh $host
        if [ $? -ne 0 ]; then
            echo "Mutual trust failed"
            return 1
        fi
    done

    # get pub_keygen_file from expect node
    # shellcheck disable=SC2068
    for host in ${expnode[@]}; do
        echo "get pub keygen file from expect $host"
        keygen_copy_exp $host
        if [ $? -ne 0 ]; then
            echo "Mutual trust failed"
            return 1
        fi
    done

    [ -f $pub_rsa ] && rm -f $pub_rsa
    return 0
}

function usage
{
    echo "Usage: `basename $0` -u userid -p passwd -l \"host1 host2...\" [-d detail message]"
    exit 1
}

while getopts 'u:p:l:hd' OPT; do
    case $OPT in
        u)
            duser="$OPTARG";;
        p)
            dpasswd="$OPTARG";;
        l)
            dips=($OPTARG);;
        d)
            detail=true;;
        h)
            usage
    esac
done

expected=false
if (cat /etc/*release | grep -Eq 'Ubuntu|Debain'); then
    if (/usr/bin/dpkg --list expect &> /dev/null); then
        expected=true
    fi
else
    if (/bin/rpm -q expect &> /dev/null); then
        expected=true
    fi
fi

if (! $expected); then
    echo "Please install expect."
    exit 1
fi

if (( ${#dips[@]} == 0 )); then
    usage
fi

echo "Setting up mutual trust for nodes: (user=$duser, passwd=$dpasswd)"
for ((i=0; i<${#dips[@]}; i++)) do
    echo "    ${dips[$i]}"
done

setup_mutual_trust
if [ $? -ne 0 ]; then
    exit 1
fi

echo -e "Finished."
