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

# Completions currently based on OpenSSH 6.0 (released on 2012-04-22).
#
# TODO: update ssh-keygen (not based on 5.9)
# TODO: sshd, ssh-keyscan, ssh-keysign

_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]'
    # for protocol version 2, this can be a comma-separated list
    '-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]:dynamic port forwarding:->dynforward' \
      '-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 smartcard device]:device:_files' \
      '-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:->macs' \
      '(-1)-N[do not execute a remote command (protocol version 2 only)]' \
      '-n[redirect stdin from /dev/null]' \
      '-O:multiplex control command:((check\:"check master process is running" exit\:"request the master to exit" forward\:"request forward without command execution" stop\:"request the master to stop accepting further multiplexing requests" cancel\:"cancel existing forwardings with -L and/or -R"))' \
      '-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]' \
      '-W[forward standard input and output to host]:stdinout forward:->hostport' \
      '-w[request tunnel device forwarding]:local_tun[\:remote_tun] (integer or "any"):' \
      '(-x -Y)-X[enable (untrusted) X11 forwarding]' \
      '(-X -Y)-x[disable X11 forwarding]' \
      '(-x -X)-Y[enable trusted X11 forwarding]' \
      '-y[send log info via syslog instead of stderr]' \
      ':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 ecdsa)" \
      "(${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
    ;;
  (ssh-copy-id)
    _arguments \
      '-i:SSH identity file:_files' \
      ':remote host name:->userhost' \
    ;;
  esac

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

    case "$lstate" in
    option)
      if compset -P '*='; then
        case "$IPREFIX" in
        *(#i)(afstokenpassing|batchmode|challengeresponseauthentication|checkhostip|clearallforwardings|compression|enablesshkeysign|exitonforwardfailure|fallbacktorsh|forward(agent|x11)|forwardx11trusted|gatewayports|gssapiauthentication|gssapidelegatecredentials|gssapitrustdns|hashknownhosts|hostbasedauthentication|identitiesonly|kbdinteractiveauthentication|(tcp|)keepalive|nohostauthenticationforlocalhost|passwordauthentication|permitlocalcommand|pubkeyauthentication|rhosts(|rsa)authentication|rsaauthentication|usersh|kerberos(authentication|tgtpassing)|useprivilegedport|visualhostkey)=*)
          _wanted values expl 'truth value' compadd yes no && ret=0
          ;;
        *(#i)addressfamily=*)
          _wanted values expl 'address family' compadd any inet inet6 && ret=0
          ;;
        *(#i)bindaddress=*)
          _wanted bind-addresses expl 'bind address' _bind_addresses && 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)compressionlevel=*)
          _values 'compression level' {1..9} && ret=0
          ;;
        *(#i)connectionattempts=*)
          _message -e 'connection attempts' && ret=0
          ;;
        *(#i)connecttimeout=*)
          _message -e 'connection timeout' && 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)controlpersist=*)
          _message -e 'timeout'
          ret=0
          _wanted values expl 'truth value' compadd yes no && ret=0
          ;;
        *(#i)escapechar=*)
          _message -e 'escape character (or `none'\'')'
          ret=0
          ;;
        *(#i)forwardx11timeout=*)
          _message -e 'timeout'
          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)hostkeyalgorithms=*)
          _values -s , 'host key algorithms' \
              'ecdsa-sha2-nistp256-cert-v01@openssh.com' \
              'ecdsa-sha2-nistp384-cert-v01@openssh.com' \
              'ecdsa-sha2-nistp521-cert-v01@openssh.com' \
              'ssh-rsa-cert-v01@openssh.com' \
              'ssh-dss-cert-v01@openssh.com' \
              'ssh-rsa-cert-v00@openssh.com' \
              'ssh-dss-cert-v00@openssh.com' \
              'ecdsa-sha2-nistp256' \
              'ecdsa-sha2-nistp384' \
              'ecdsa-sha2-nistp521' \
              'ssh-rsa' \
              'ssh-dss' && ret=0
          ;;
        *(#i)identityfile=*)
          _description files expl 'SSH identity file'
          _files "$expl[@]" && ret=0
          ;;
        *(#i)ipqos=*)
          local descr
          if [[ $PREFIX = *\ *\ * ]]; then return 1; fi
          if compset -P '* '; then
            descr='QoS for non-interactive sessions'
          else
            descr='QoS [for interactive sessions if second value given, separated by white space]'
          fi
          _values $descr 'af11' 'af12' 'af13' 'af14' 'af22' \
              'af23' 'af31' 'af32' 'af33' 'af41' 'af42' 'af43' \
              'cs0' 'cs1' 'cs2' 'cs3' 'cs4' 'cs5' 'cs6' 'cs7' 'ef' \
              'lowdelay' 'throughput' 'reliability' && ret=0
          ;;
        *(#i)(local|remote)forward=*)
          state=forward
          ;;
        *(#i)dynamicforward=*)
          state=dynforward
          ;;
        *(#i)kbdinteractivedevices=*)
          _values -s , 'keyboard-interactive authentication methods' \
              'bsdauth' 'pam' 'skey' && ret=0
          ;;
        *(#i)kexalgorithms=*)
          _values -s , 'KEX algorithms' \
              ecdh-sha2-nistp256 ecdh-sha2-nistp384 ecdh-sha2-nistp521 \
              diffie-hellman-group-exchange-sha256 \
              diffie-hellman-group-exchange-sha1 \
              diffie-hellman-group14-sha1 diffie-hellman-group1-sha1 && ret=0
          ;;
        *(#i)localcommand=*)
          _description commands expl 'run command locally after connecting'
          _command_names && ret=0
          ;;
        *(#i)loglevel=*)
          _values 'log level' QUIET FATAL ERROR INFO VERBOSE\
              DEBUG DEBUG1 DEBUG2 DEBUG3 && ret=0
          ;;
        *(#i)macs=*)
          state=macs
          ;;
        *(#i)numberofpasswordprompts=*)
          _message -e 'number of password prompts'
          ret=0
          ;;
        *(#i)pkcs11provider=*)
          _description files expl 'PKCS#11 shared library'
          _files -g '*.so' "$expl[@]" && ret=0
          ;;
        *(#i)port=*)
          _message -e 'port number on remote host'
          ret=0
          ;;
        *(#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)rekeylimit=*)
          _message -e 'maximum number of bytes transmitted before renegotiating session key'
          ret=0
          ;;
        *(#i)requesttty=*)
          _values 'request a pseudo-tty' \
              'no[never request a TTY]' \
              'yes[always request a TTY when stdin is a TTY]' \
              'force[always request a TTY]' \
              'auto[request a TTY when opening a login session]' && ret=0
          ;;
        *(#i)sendenv=*)
          _wanted envs expl 'environment variable' _parameters -g 'scalar*export*' && ret=0
          ;;
        *(#i)serveralivecountmax=*)
          _message -e 'number of alive messages without replies before disconnecting'
          ret=0
          ;;
        *(#i)serveraliveinterval=*)
          _message -e 'timeout in seconds since last data was received to send alive message'
          ret=0
          ;;
        *(#i)(stricthostkeychecking|verifyhostkeydns)=*)
          _wanted values expl 'checking type' compadd yes no ask && ret=0
          ;;
        *(#i)tunnel=*)
          _values 'request device forwarding' \
              'yes' \
              'point-to-point' \
              'ethernet' \
              'no' && ret=0
          ;;
        *(#i)tunneldevice=*)
          _message -e 'local_tun[:remote_tun] (integer or "any")'
          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}' -S '=' - \
                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 \
                GSSAPITrustDns \
                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)
      local port=false host=false listen=false bind=false
      if compset -P 1 '*:'; then
        if [[ $IPREFIX != (*=|)<-65535>: ]]; then
          if compset -P 1 '*:'; then
            if compset -P '*:'; then
              port=true
            else
              host=true
            fi
          else
            listen=true
            ret=0
          fi
        else
          if compset -P '*:'; then
            port=true
          else
            host=true
          fi
        fi
      else
        listen=true
        bind=true
      fi
      $port && { _message -e port-numbers 'port number'; ret=0 }
      $listen && { _message -e port-numbers 'listen-port number'; ret=0 }
      $host && { _wanted hosts expl host _ssh_hosts -S: && ret=0 }
      $bind && { _wanted bind-addresses expl bind-address _bind_addresses -S: && ret=0 }
      return ret
      ;;
    dynforward)
      _message -e port-numbers 'listen-port number'
      if ! compset -P '*:'; then
        _wanted bind-addresses expl bind-address _bind_addresses -qS:
      fi
      return 0
      ;;
    hostport)
      if compset -P '*:'; then
        _message -e port-numbers 'port number'
        ret=0
      else
        _wanted hosts expl host _ssh_hosts -S: && ret=0
      fi
      return ret
      ;;
    macs)
      _values -s , 'MAC algorithms' hmac-md5 hmac-sha1 umac-64@openssh.com \
          hmac-ripemd160 hmac-sha1-96 hmac-md5-96 hmac-sha2-256 \
          hmac-sha2-256-96 hmac-sha2-512 hmac-sha2-512-96 && ret=0
      ;;
    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 -- ssh ${(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 -- ssh && 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

  return ret
}

_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 "$@"