#compdef ssh slogin=ssh scp ssh-add ssh-agent ssh-keygen sftp

# Completions currently based on OpenSSH 5.9 (released on 2011-09-06).
#
# TODO: update ssh-keygen (not based on 5.9)
# TODO: sshd, ssh-keyscan, ssh-keysign


_remote_files () {
  # There should be coloring based on all the different ls -F classifiers.
  local expl rempat remfiles remdispf remdispd args suf ret=1

  if zstyle -T ":completion:${curcontext}:files" remote-access; then
    zparseopts -D -E -a args p: 1 2 4 6 F:
    if [[ -z $QIPREFIX ]]
    then rempat="${PREFIX%%[^./][^/]#}\*"
    else rempat="${(q)PREFIX%%[^./][^/]#}\*"
    fi
    remfiles=(${(M)${(f)"$(_call_program files ssh -o BatchMode=yes $args -a -x ${IPREFIX%:} ls -d1FL -- "$rempat" 2>/dev/null)"}%%[^/]#(|/)})
    compset -P '*/'
    compset -S '/*' || suf='remote file'

    remdispf=(${remfiles:#*/})
    remdispd=(${(M)remfiles:#*/})

    _tags files
    while _tags; do
      while _next_label files expl ${suf:-remote directory}; do
        [[ -n $suf ]] && compadd "$@" "$expl[@]" -d remdispf \
	    ${(q)remdispf%[*=@|]} && ret=0 
	compadd ${suf:+-S/} "$@" "$expl[@]" -d remdispd \
	    ${(q)remdispd%/} && ret=0
      done
      (( ret )) || return 0
    done
    return ret
  else
    _message -e remote-files 'remote file'
  fi
}

_ssh () {
  local curcontext="$curcontext" state line expl common tmp cmds suf ret=1
  typeset -A opt_args

  common=(
    '(-2)-1[forces ssh to try protocol version 1 only]'
    '(-1)-2[forces ssh to try protocol version 2 only]'
    '(-6)-4[forces ssh to use IPv4 addresses only]'
    '(-4)-6[forces ssh to use IPv6 addresses only]'
    '-C[compress data]'
    '-c+[select encryption cipher]:encryption cipher:(idea des 3des blowfish arcfour tss none)'
    '-F+[specify alternate config file]:config file:_files'
    '-i+[select identity file]:SSH identity file:_files'
    '*-o+[specify extra options]:option string:->option'
  )
  common_transfer=(
    '-l[limit used bandwidth]:bandwidth in KiB/s:'
    '-P+[specify port on remote host]:port number on remote host'
    '-p[preserve modification times, access times and modes]'
    '-q[disable progress meter and warnings]'
    '-r[recursively copy directories (follows symbolic links)]'
    '-S+[specify ssh program]:path to ssh:_command_names -e' \
    '-v[verbose mode]'
  )

  case "$service" in
  ssh)
    _arguments -C -s \
      '(-a)-A[enables forwarding of the authentication agent connection]' \
      '(-A)-a[disable forwarding of authentication agent connection]' \
      '(-P)-b+[specify interface to transmit on]:bind address:_bind_addresses' \
      '-D+[specify a dynamic port forwarding]:[bind-address]\:port' \
      '-e+[set escape character]:escape character (or `none'"'"'):' \
      '(-n)-f[go to background]' \
      '-g[allow remote hosts to connect to local forwarded ports]' \
      '-I+[specify the PKCS#11 shared library to use]' \
      '-K[enable GSSAPI-based authentication and forwarding]' \
      '-k[disable forwarding of GSSAPI credentials]' \
      '*-L[specify local port forwarding]:local port forwarding:->forward' \
      '-l+[specify login name]:login name:_ssh_users' \
      '-M[master mode for connection sharing]' \
      '(-1)-m+[specify mac algorithms]:mac spec' \
      '(-1)-N[do not execute a remote command (protocol version 2 only)]' \
      '-n[redirect stdin from /dev/null]' \
      '-O[control active connection multiplexing master process]:multiplex control command:((
          check\:"check that the master process is running"
          forward\:"request forwardings without command execution"
          cancel\:"cancel forwardings"
          exit\:"request the master to exit"
          stop\:"request the master to stop accepting further multiplexing requests"))' \
      '-P[use non privileged port]' \
      '-p+[specify port on remote host]:port number on remote host' \
      '(-v)*-q[quiet operation]' \
      '*-R[specify remote port forwarding]:remote port forwarding:->forward' \
      '-S+[specify location of control socket for connection sharing]:path to control socket:_files' \
      '(-1)-s[invoke subsystem]' \
      '(-1 -t)-T[disable pseudo-tty allocation (protocol version 2 only)]' \
      '(-T)-t[force pseudo-tty allocation]' \
      '-V[show version number]' \
      '(-q)*-v[verbose mode]' \
      '(-N)-W[forward standard input/output over host:port (protocol version 2 only)]:host\:port' \
      '-w[request tunnel device forwarding with the specified tun devices]:local_tun[\:remote_tun]' \
      '(-x -Y)-X[enable (untrusted) X11 forwarding]' \
      '(-X -Y)-x[disable X11 forwarding]' \
      '(-x -X)-Y[enable trusted X11 forwarding]' \
      '-y[send log information using the syslog module]' \
      ':remote host name:->userhost' \
      '*::args:->command' "$common[@]" && ret=0
    ;;
  scp)
    _arguments -C -s \
      '-3[copy through local host, not directly between the remote hosts]' \
      '-B[batch mode (don'"'"'t ask for passphrases)]' \
      '*:file:->file' "$common[@]" "$common_transfer[@]" && ret=0
    ;;
  ssh-add)
    _arguments -s \
      '-c[identity is subject to confirmation via SSH_ASKPASS]' \
      '-D[delete all identities]' \
      '-d[remove identity]' \
      '-e[remove keys provided by the PKCS#11 shared library]:library:' \
      '-k[load plain private keys only and skip certificates]' \
      '-L[lists public key parameters of all identities in the agent]'\
      '-l[list all identities]' \
      '-s[add keys provided by the PKCS#11 shared library]:library:' \
      '-t[set maximum lifetime for identity]:maximum lifetime (in seconds or time format):' \
      '-X[unlock the agent]' \
      '-x[lock the agent with a password]' \
      '*:SSH identity file:_files'
    return
    ;;
  ssh-agent)
    _arguments -s \
      '(-k)-a[UNIX-domain socket to bind agent to]:UNIX-domain socket:_files' \
      '(-k -s)-c[force csh-style shell]' \
      '(-k)-d[debug mode]' \
      '-k[kill current agent]' \
      '(-k -c)-s[force sh-style shell]' \
      '-t[set default maximum lifetime for identities]:maximum lifetime (in seconds or time format):' \
      '*::command: _normal'
    return
    ;;
  ssh-keygen)
    cmds=( -p -i -e -y -c -l -B -D -U )
    _arguments \
      '-q[silence ssh-keygen]' \
      "($cmds -P)-b[specify number of bits in key]:bits in key" \
      "($cmds -P)-t[specify the type of the key to create]:key type:(rsa1 rsa dsa)" \
      "(${cmds#-p })-N[provide new passphrase]:new passphrase" \
      "($cmds -b -t)-C[provide new comment]:new comment" \
      '(-D)-f[key file]:key file:_files' \
      '('${(j. .)cmds:#-[pc]}' -t -b)-P[provide old passphrase]:old passphrase' \
      "($cmds -q -b -t -C)-p[change passphrase of private key file]" \
      "($cmds -q -b -t -N -C -P)-i[import key to OpenSSH format]" \
      "($cmds -q -b -t -N -C -P)-e[export key to SECSH file format]" \
      "($cmds -q -b -t -N -C -P)-y[get public key from private key]" \
      "($cmds -q -b -t -N)-c[change comment in private and public key files]" \
      "($cmds -q -b -t -N -C -P)-l[show fingerprint of key file]" \
      "($cmds -q -b -t -N -C -P)-B[show the bubblebabble digest of key]" \
      "($cmds -q -b -t -N -C -P -f)-D[download key stored in smartcard reader]:reader" \
      "($cmds -q -b -t -N -C -P)-U[upload key to smartcard reader]:reader"
    return
    ;;
  sftp)
    _arguments -C -s \
      '-B+[specify buffer size]:buffer size in bytes (default\: 32768):' \
      '-b+[specify batch file to read]:batch file:_files' \
      '-D[connect directly to a local sftp server]:sftp server path:' \
      '-R[specify number of outstanding requests]:number of requests (default\: 64):' \
      '-s[SSH2 subsystem or path to sftp server on the remote host]' \
      '1:file:->rfile' '*:file:->file' "$common[@]" "$common_transfer[@]" && ret=0
    ;;
  esac

  while [[ -n "$state" ]]; do
    lstate="$state"
    state=''

    case "$lstate" in
    option)
      if compset -P '*[= ]'; then
        case "$IPREFIX" in
        *(#i)(afstokenpassing|batchmode|compression|fallbacktorsh|forward(agent|x11)|keepalive|passwordauthentication|rhosts(|rsa)authentication|rsaauthentication|usersh|kerberos(authetication|tgtparsing)|useprivileged)*)
	  _wanted values expl 'truth value' compadd yes no && ret=0
          ;;
        *(#i)ciphers*)
          _values -s , 'encryption cipher' \
	      '3des-cbc' \
	      'aes128-cbc' \
	      'aes192-cbc' \
	      'aes256-cbc' \
	      'aes128-ctr' \
	      'aes192-ctr' \
	      'aes256-ctr' \
	      'arcfour128' \
	      'arcfour256' \
	      'arcfour' \
	      'blowfish-cbc' \
	      'cast128-cbc' \
	      \
	      'rijndael128-cbc' \
	      'rijndael192-cbc' \
	      'rijndael256-cbc' \
	      'rijndael-cbc@lysator.liu.se' \
	      && ret=0
          ;;
        *(#i)cipher*)
	  _wanted values expl 'encryption cipher (protocol version 1)' \
              compadd blowfish 3des des idea arcfour tss none && ret=0
          ;;
	*(#i)controlmaster*)
	  _wanted values expl 'truthish value' compadd yes no auto autoask && ret=0
	  ;;
	*(#i)controlpath*)
          _description files expl 'path to control socket'
          _files "$expl[@]" && ret=0
	  ;;
        *(#i)globalknownhostsfile*)
          _description files expl 'global file with known hosts'
          _files "$expl[@]" && ret=0
          ;;
        *(#i)hostname*)
	  _wanted hosts expl 'real host name to log into' _ssh_hosts && ret=0
          ;;
        *(#i)identityfile*)
          _description files expl 'SSH identity file'
          _files "$expl[@]" && ret=0
          ;;
        *(#i)(local|remote)forward*)
          state=forward
          ;;
        *(#i)preferredauthentications*)
          _values -s , 'authentication method' gssapi-with-mic \
              hostbased publickey keyboard-interactive password && ret=0
          ;;
        *(#i)protocol*)
          _values -s , 'protocol version' \
	      '1' \
	      '2' && ret=0
	  ;;
        *(#i)proxycommand*)
          compset -q
          shift 1 words
          (( CURRENT-- ))
          _normal && ret=0
          ;;
        *(#i)stricthostkeychecking*)
          _wanted values expl 'checking type' compadd yes no ask && ret=0
          ;;
        *(#i)userknownhostsfile*)
          _description files expl 'user file with known hosts'
          _files "$expl[@]" && ret=0
          ;;
        *(#i)user*)
	  _wanted users expl 'user to log in as' _ssh_users && ret=0
          ;;
        *(#i)xauthlocation*)
          _description files expl 'xauth program'
          _files "$expl[@]" -g '*(-*)' && ret=0
          ;;
        esac
      else
        # old options are after the empty "\"-line
        _wanted values expl 'configure file option' \
            compadd -M 'm:{a-z}={A-Z}' -qS '=' - \
                AddressFamily \
                BatchMode \
                BindAddress \
                ChallengeResponseAuthentication \
                CheckHostIP \
                Cipher \
                Ciphers \
                ClearAllForwardings \
                Compression \
                CompressionLevel \
                ConnectionAttempts \
                ConnectTimeout \
                ControlMaster \
                ControlPath \
                ControlPersist \
                DynamicForward \
                EnableSSHKeysign \
                EscapeChar \
                ExitOnForwardFailure \
                ForwardAgent \
                ForwardX11 \
                ForwardX11Timeout \
                ForwardX11Trusted \
                GatewayPorts \
                GlobalKnownHostsFile \
                GSSAPIAuthentication \
                GSSAPIDelegateCredentials \
                HashKnownHosts \
                Host \
                HostbasedAuthentication \
                HostKeyAlgorithms \
                HostKeyAlias \
                HostName \
                IdentitiesOnly \
                IdentityFile \
                IPQoS \
                KbdInteractiveAuthentication \
                KbdInteractiveDevices \
                KexAlgorithms \
                LocalCommand \
                LocalForward \
                LogLevel \
                MACs \
                NoHostAuthenticationForLocalhost \
                NumberOfPasswordPrompts \
                PasswordAuthentication \
                PermitLocalCommand \
                PKCS11Provider \
                Port \
                PreferredAuthentications \
                Protocol \
                ProxyCommand \
                PubkeyAuthentication \
                RekeyLimit \
                RemoteForward \
                RequestTTY \
                RhostsRSAAuthentication \
                RSAAuthentication \
                SendEnv \
                ServerAliveCountMax \
                ServerAliveInterval \
                StrictHostKeyChecking \
                TCPKeepAlive \
                Tunnel \
                TunnelDevice \
                UsePrivilegedPort \
                User \
                UserKnownHostsFile \
                VerifyHostKeyDNS \
                VisualHostKey \
                XAuthLocation \
                \
                AFSTokenPassing \
                FallBackToRsh \
                KeepAlive \
                KerberosAuthentication \
                KerberosTgtPassing \
                PreferredAuthentications \
                ProtocolKeepAlives \
                RhostsAuthentication \
                SetupTimeOut \
                SmartcardDevice \
                UseRsh \
                && ret=0
      fi
      ;;
    forward)
      if compset -P 1 '*:'; then
        if compset -P '*:'; then
          _message -e port-numbers 'port number'
        else
	  _wanted hosts expl host _ssh_hosts -qS:
        fi
      else
        _message -e port-numbers 'listen-port number'
      fi
      return
      ;;
    command)
      shift 1 words
      (( CURRENT-- ))
      _normal
      return
      ;;
    userhost)
      if compset -P '*@'; then
	_wanted hosts expl 'remote host name' _ssh_hosts && ret=0
      elif compset -S '@*'; then
        _wanted users expl 'login name' _ssh_users -S '' && ret=0
      else
        if (( $+opt_args[-l] )); then
	  tmp=()
	else
	  tmp=( 'users:login name:_ssh_users -qS@' )
	fi
	_alternative \
	    'hosts:remote host name:_ssh_hosts' \
	    "$tmp[@]" && ret=0
      fi
      ;;
    file)
      if compset -P '*:'; then
        _remote_files ${(kv)~opt_args[(I)-[FP1246]]/-P/-p} && ret=0
      elif compset -P '*@'; then
        suf=( -S '' )
        compset -S ':*' || suf=( -r: -S: )
        _wanted hosts expl 'remote host name' _ssh_hosts $suf && ret=0
      else
        _alternative \
	    'files:: _files' \
	    'hosts:remote host name:_ssh_hosts -r: -S:' \
	    'users:user:_ssh_users -qS@' && ret=0
      fi
      ;;
    rfile)
      if compset -P '*:'; then
        _remote_files && ret=0
      elif compset -P '*@'; then
        _wanted hosts expl host _ssh_hosts -r: -S: && ret=0
      else
        _alternative \
	    'hosts:remote host name:_ssh_hosts -r: -S:' \
	    'users:user:_ssh_users -qS@' && ret=0
      fi
      ;;
    esac
  done
}

_ssh_users () {
  _combination -s '[:@]' my-accounts users-hosts users "$@"
}

_ssh_hosts () {
  local -a config_hosts
  local config
  integer ind

  # If users-hosts matches, we shouldn't complete anything else.
  if [[ "$IPREFIX" == *@ ]]; then
    _combination -s '[:@]' my-accounts users-hosts "users=${IPREFIX/@}" hosts "$@" && return
  else
    _combination -s '[:@]' my-accounts users-hosts \
      ${opt_args[-l]:+"users=${opt_args[-l]:q}"} hosts "$@" && return
  fi
  if (( ind = ${words[(I)-F]} )); then
    config=${~words[ind+1]}
  else
    config="$HOME/.ssh/config"
  fi
  if [[ -r $config ]]; then
    local IFS=$'\t ' key hosts host
    while read key hosts; do
      if [[ "$key" == (#i)host ]]; then
	 for host in ${(z)hosts}; do
	    case $host in
	    (*[*?]*) ;;
	    (*) config_hosts+=("$host") ;;
	    esac
	 done
      fi
    done < "$config"
    if (( ${#config_hosts} )); then
      _wanted hosts expl 'remote host name' \
	compadd -M 'm:{a-zA-Z}={A-Za-z} r:|.=* r:|=*' "$@" $config_hosts
    fi
  fi
}

_ssh "$@"
