# This function performs host completion based on ssh's known_hosts files,
# defaulting to standard host completion if they don't exist.


comp_include _get_cword


_known_hosts()
{
       local cur curd ocur user suffix aliases global_kh user_kh hosts i host
       local -a kh khd config

    COMPREPLY=()
    cur=`_get_cword`
    ocur=$cur

    [ "$1" = -a ] || [ "$2" = -a ] && aliases='yes'
    [ "$1" = -c ] || [ "$2" = -c ] && suffix=':'
    [[ $cur == *@* ]] && user=${cur%@*}@ && cur=${cur#*@}
    kh=()

    # ssh config files
    [ -r /etc/ssh/ssh_config ] &&
      config=( ${config[@]} /etc/ssh/ssh_config )
    [ -r ~/.ssh/config ] &&
      config=( ${config[@]} ~/.ssh/config )
    [ -r ~/.ssh2/config ] &&
      config=( ${config[@]} ~/.ssh2/config )

    if [ ${#config[@]} -gt 0 ]; then
        # expand path (if present) to global known hosts file
        global_kh=$( eval echo $( sed -ne 's/^[Gg][Ll][Oo][Bb][Aa][Ll][Kk][Nn][Oo][Ww][Nn][Hh][Oo][Ss][Tt][Ss][Ff][Ii][Ll][Ee]['"$'\t '"']*\(.*\)$/\1/p' ${config[@]} ) )
        # expand path (if present) to user known hosts file
        user_kh=$( eval echo $( sed -ne 's/^[Uu][Ss][Ee][Rr][Kk][Nn][Oo][Ww][Nn][Hh][Oo][Ss][Tt][Ss][Ff][Ii][Ll][Ee]['"$'\t '"']*\(.*\)$/\1/p' ${config[@]} ) )
    fi

    # choose which global known hosts file to use
    if [ -r "$global_kh" ]; then
        kh=( "$global_kh" )
    else
        [ -r /etc/ssh/ssh_known_hosts ] &&
          kh=( ${kh[@]} /etc/ssh/ssh_known_hosts )
        [ -r /etc/ssh/ssh_known_hosts2 ] &&
          kh=( ${kh[@]} /etc/ssh/ssh_known_hosts2 )
        [ -r /etc/known_hosts ] &&
          kh=( ${kh[@]} /etc/known_hosts )
        [ -r /etc/known_hosts2 ] &&
          kh=( ${kh[@]} /etc/known_hosts2 )
        [ -d /etc/ssh2/knownhosts ] &&
          khd=( ${khd[@]} /etc/ssh2/knownhosts/*pub )
    fi

    # choose which user known hosts file to use
    if [ -r "$user_kh" ]; then
        kh=( ${kh[@]} "$user_kh" )
    else
        [ -r ~/.ssh/known_hosts ] &&
          kh=( ${kh[@]} ~/.ssh/known_hosts )
        [ -r ~/.ssh/known_hosts2 ] &&
          kh=( ${kh[@]} ~/.ssh/known_hosts2 )
        [ -d ~/.ssh2/hostkeys ] &&
          khd=( ${khd[@]} ~/.ssh2/hostkeys/*pub )
    fi

    # If we have known_hosts files to use
    if [ ${#kh[@]} -gt 0 -o ${#khd[@]} -gt 0 ]; then
        # Escape slashes and dots in paths for awk
        cur=${cur//\//\\\/}
        cur=${cur//\./\\\.}
        curd=$cur

        if [[ "$cur" == [0-9]*.* ]]; then
        # Digits followed by a dot - just search for that
        cur="^$cur.*"
        elif [[ "$cur" == [0-9]* ]]; then
        # Digits followed by no dot - search for digits followed
        # by a dot
        cur="^$cur.*\."
        elif [ -z "$cur" ]; then
        # A blank - search for a dot or an alpha character
        cur="[a-z.]"
        else
        cur="^$cur"
        fi

        if [ ${#kh[@]} -gt 0 ]; then

        # FS needs to look for a comma separated list
        COMPREPLY=( $( awk 'BEGIN {FS=","}
                {for (i=1; i<=2; ++i) { \
                       gsub(" .*$", "", $i); \
                       if ($i ~ /'$cur'/) {print $i} \
                }}' ${kh[@]} 2>/dev/null ) )
        fi
        if [ ${#khd[@]} -gt 0 ]; then
        # Needs to look for files called
        # .../.ssh2/key_22_<hostname>.pub
        # dont fork any processes, because in a cluster environment, 
        # there can be hundreds of hostkeys
        for i in ${khd[@]} ; do
            if [[ "$i" == *key_22_$curd*.pub ]] && [ -r "$i" ] ; then
            host=${i/#*key_22_/}
            host=${host/%.pub/}
            COMPREPLY=( ${COMPREPLY[@]} $host )
            fi
        done
        fi
        # append any available aliases from config files
        if [ ${#config[@]} -gt 0 ] && [ -n "$aliases" ]; then
        hosts=$( compgen -W "$( sed -ne 's/^[Hh][Oo][Ss][Tt]['"$'\t '"']*\([^*?]*\)$/\1/p' ${config[@]} )" -- $ocur )
        COMPREPLY=( ${COMPREPLY[@]} $hosts )
        fi

        # apply suffix
        for (( i=0; i < ${#COMPREPLY[@]}; i++ )); do
        COMPREPLY[i]=$user${COMPREPLY[i]}$suffix
        done
    else
        # Just do normal hostname completion
        COMPREPLY=( $( compgen -A hostname -S "$suffix" -- $cur ) )
    fi

    return 0
}
_known_hosts
