#!/bin/bash

# Private and public key file
HOME_DIR=`cd -P ~/; pwd -P`
PRIVATE_ID="${HOME_DIR}/.ssh/id_rsa"
PUBLIC_ID="${HOME_DIR}/.ssh/id_rsa.pub"

# Default user and password
type="host"
user="root"
password="huawei"
new=0

Usage() {
    echo "Usage: sh config-ssh.sh -h <ip>[,...] -t <type> -u <user> -p <password> [-n]"
    echo "       -h: Host ip"
    echo "       -t: Device type (host/device, Default: host)"
    echo "       -u: Username (Default: root)"
    echo "       -p: Password (Default: huawei)"
    echo "       -n: Generate a new rsa code"
    exit 1
}

# Check input parameter
if [ $# -eq 0 ]; then
    Usage
fi

# Parse parameter
while getopts ":h:t:u:p:n" OPTIONS; do
    case ${OPTIONS} in
        h) ipaddr="$OPTARG";;
        t) type="$OPTARG";;
        u) user="$OPTARG";;
        p) password="$OPTARG";;
        n) new=1;;
        *) Usage;;
    esac
done

InstallRPM() {
    OS_TYPE=`grep "^ID=" /etc/os-release | awk -F'=' '{print $2}'`
    if [ "${OS_TYPE}" = "ubuntu" ]; then
        sudo apt-get install -y tcl > /dev/null 2>&1
        sudo apt-get install -y expect > /dev/null 2>&1
    else
        OS_TYPE=`echo ${OS_TYPE} | awk -F'"' '{print $2}'`
        if [ "${OS_TYPE}" = "EulerOS" ] || [ "${OS_TYPE}" = "centos" ]; then
            sudo yum install -y tcl > /dev/null 2>&1
            sudo yum install -y expect > /dev/null 2>&1
        elif [ "${OS_TYPE}" = "sles" ]; then
            sudo zypper install -y tcl > /dev/null 2>&1
            sudo zypper install -y sel > /dev/null 2>&1
        fi
    fi
}

SendIdPub() {
    local file=$1
    local dest=$2
    local user_dst=$3
    local device_type=$4
    if [ "${device_type}" != "device" ]; then
        expect -c "spawn ssh-copy-id  -i ${file} ${user_dst}@${dest}
                   expect {
                           \"Are you sure you want to continue connecting (yes/no)?\" {
                                   send \"yes\r\"
                                   expect -re \"\[P|p]assword:\"
                                   send \"${password}\r\"
                           }
                           ## already has public key in ~/.ssh/known_hosts
                           -re \"\[P|p]assword:\" {
                                   send \"${password}\r\"
                                   }
                           ## Do not need input password. Becauese of ssh
                           eof {
                                   catch wait result
                                   #send_user  [lindex \$result 3]
                                   exit [lindex \$result 3]
                           }
                   }
                   ### We have input password,and the command may have been execute,except password is wrong or connctione is broken.
                       expect {
                           ## check exit status of the proccess of ssh
                           eof {
                                                            catch wait result
                                   exit [lindex \$result 3]
                           }
                           ## Password is wrong!
                           -re \"\[P|p]assword:\" {
                                   send_user \"invalid password or account. \$expect_out(buffer)\n\"
                                   exit 13
                           }
                   }"
    else
        local id_rsa_pub=`cat ${file}`
        local id_rsa_pub_name=`echo ${id_rsa_pub}|awk '{print $3}'`
        ssh-keygen -f "${HOME_DIR}/.ssh/known_hosts" -R ${dest}
        expect -c "spawn ssh ${user}@${dest} -y \"/usr/sbin/mkdir -pv ~/.ssh; sed -i /$id_rsa_pub_name/d ~/.ssh/authorized_keys; echo $id_rsa_pub >> ~/.ssh/authorized_keys\"
                   expect {
                           \"Are you sure you want to continue connecting (yes/no)?\" {
                                   send \"yes\r\"
                                   expect -re \"\[P|p]assword:\"
                                   send \"${password}\r\"
                           }
                           ## already has public key in ~/.ssh/known_hosts
                           -re \"\[P|p]assword:\" {
                                   send \"${password}\r\"
                                   }
                           ## Do not need input password. Becauese of ssh
                           eof {
                                   catch wait result
                                   #send_user  [lindex \$result 3]
                                   exit [lindex \$result 3]
                           }
                   }
                   ### We have input password,and the command may have been execute,except password is wrong or connctione is broken.
                       expect {
                           ## check exit status of the proccess of ssh
                           eof {
                                   catch wait result
                                   exit [lindex \$result 3]
                           }
                           ## Password is wrong!
                           -re \"\[P|p]assword:\" {
                                   send_user \"invalid password or account. \$expect_out(buffer)\n\"
                                   exit 13
                           }
                   }"
    fi
    return $?
}

main () {
    # Install dependent package
    InstallRPM
    
    # Get array of host IP
    IP_LIST=`echo ${ipaddr} | sed s/,/\ /g`
    #echo "Host IP: ${IP_LIST}"
    if [ -z "${IP_LIST}" ]; then
        Usage
    fi
    
    # Check network of host IP
    for ip in ${IP_LIST}; do
        if [ "${type}" != "device" ]; then
            ping ${ip} -c 1 1>/dev/null
            if [ "$?" -ne 0 ]; then
                echo "ERROR: ${ip} can not be reached."
                exit 1
            fi
        fi
    done

    # Generate new rsa
    if [ -f "${PUBLIC_ID}" -a "${new}" = "0" ]; then
        :
    else
        mkdir -p /root/.ssh > /dev/null 2>&1
        rm -f ${PRIVATE_ID} ${PUBLIC_ID} > /dev/null 2>&1
        ssh-keygen -t rsa -P '' -f ${PRIVATE_ID} > ~/config-ssh.log 2>&1
    fi
    
    # Configure ssh
    for ip in ${IP_LIST}; do
        sed -i /${ip}/d ~/.ssh/known_hosts > ~/config-ssh.log 2>&1
        SendIdPub "${PUBLIC_ID}" "${ip}" "${user}" "${type}" > ~/config-ssh.log 2>&1
        if [ $? -eq 0 ]; then
            echo "Configure ${user}@${ip} Done."
        else
            echo "Configure ${user}@${ip} Fail."
        fi
    done
}

main $@
