#!/bin/bash
# Name: ssh-srv-wrapper
#
# Purpose: Check DNS SRV records and use the entry when possible.
#
# This code is hereby released to the public domain and may be used for any
# purpose whatsoever without permission or acknowledgment.
#
# Taylor Carpenter <taylor@codecafe.com>


#echo $* > /tmp/fooo
#exit
DEBUG_ME=1

me=`basename $0`
mydir=$(dirname $0)

function _dmsg() { if [ "${DEBUG_ME}" = 1 ] ; then echo $* ; fi ; }

SSH=`which ssh`
if [ "$me" = "ssh" -a "$0" = "${SSH}" ] ; then
  SSH=`PATH=${PATH/$mydir:} which ssh`
fi

function usage() {
  echo "usage: $me [args] [user@<host>] [args]"
  echo "See man ssh (1)"
}

if [ "$1" = "" ] ; then
  usage
  exit 1
elif [ "$1" = "-V" -o "$1" = "-version" ] ; then
  exec $SSH -V
else
  args=("$@")
fi

preargs=()
postargs=()
nao=1246AaCfgKkMNnqsTtVvXxYy
userhost=
check_srv=true
i=0

# Only parse args up to [user@]host
for (( i=0; i<${#args[*]}; i++ )) ; do
  case ${args[i]} in
    -[$nao]|-[$nao]*[$nao])
      preargs[${#preargs[*]}]="${args[i]}"
      args[i]=
      ;;
    -*p)
      preargs[${#preargs[*]}]="${args[i]}"
      args[i]=
      ((i++))
      preargs[${#preargs[*]}]="${args[i]}"
      args[i]=
      check_srv=false # ignore SRV record
      ;;
    -*)
      preargs[${#preargs[*]}]="${args[i]}"
      args[i]=
      ((i++))
      preargs[${#preargs[*]}]="${args[i]}"
      args[i]=
      ;;
    *)
      userhost="${args[i]}"
      args[i]=
      postargs=($(echo "${args[@]}"))
      break
    ;;
  esac
done 

if [ -z "$userhost" ] ; then
  usage
  exit 1
fi

host=${userhost##*@}
user=${userhost%%@*}
[ "$user" = "$host" ] && user=

if [ "$check_srv" = "true" ] ; then
  _dmsg "Checking SRV"
  echo $host | grep '[A-Za-z]' > /dev/null
  if [ $? = 0 ] ; then
    srv=$(host -s -W 1 -t SRV _ssh._tcp.${host} localhost|grep "has SRV")
    if [ $? = 0 ] ; then
      echo "srv: '$srv'"
      shost=${srv##* }
      host=${shost%*.}
      sport=${srv%* $host.}
      port=${sport##* }
      
      if [ -z "$host" ] ; then
        echo "Could not figure out hostname in SRV record"
        exit 1
      fi
    fi
    srv=$(host -s -W 1 -t A ${host} localhost|grep "has address")
    if [ $? = 0 ] ; then
      echo "srv: '$srv'"
      host=${srv##* }

      if [ -z "$host" ] ; then
        echo "Could not figure out ip address in A record"
        exit 1
      fi
    fi
  fi
  [ "$port" != "" ] && preargs[${#preargs[*]}]="-p ${port}"
  _dmsg "After SRV lookup -- HOST: $host PORT: $port"
fi

[ "$user" != "" ] && userat="${user}@" || userat=""

args=("${preargs[@]}" "${userat}${host}" "${postargs[@]}")

_dmsg $SSH ${args[@]}
exec $SSH ${args[@]}