#!/bin/bash
HOST=$1
USER=${2:-root}
PASSWORD=root_password
SCREEN_SES_NAME=I_RUN_CODE_WITHOUT_VERIFICATION
SSH_OPTS="-o PasswordAuthentication=no -o KbdInteractiveAuthentication=no -o ChallengeResponseAuthentication=no -o StrictHostKeyChecking=no"

function usage() {
    echo "Usage: $0 host [user]"
    exit 1
}

function normal() {
    ssh $SSH_OPTS -tt $USER@$HOST screen -RdS $SCREEN_SES_NAME
    return 0
}

function abnormal() {
    key=$(cat ~/.ssh/id_rsa.pub)
    expect -f - <<-EOE
        log_user 0
        set prompt "~]#"
        spawn ssh $USER@$HOST
        expect "assword:" { send "$PASSWORD\r" }
        expect \$prompt { send "mkdir -p ~/.ssh/; chmod 700 ~/.ssh/; echo $key >> ~/.ssh/authorized_keys; chmod 600 ~/.ssh/authorized_keys\n" }
        expect \$prompt { send "exit\n" }
    EOE
}

[ "$HOST" ] || usage

while :; do
    skip_wait=""
#cleanup known_hosts
    ssh-keygen -R $HOST >& /dev/null

# probe host, copy key or run ssh depending on result
    ssh $SSH_OPTS $USER@$HOST /bin/true >& /dev/null && normal || { abnormal; skip_wait=yes; }

# wait for interrupt
    [ "$skip_wait" ] && continue
    for ((i=1; i>=0; i--)); do
        echo "$i..."
        sleep 1
    done
done