SSH + iTerm + SRV Wrapper

#!/bin/bash

##
## Purpose: This script does a few different things before it opens an SSH session. First, it
## uses the magical DNS SRV records to make sure I'm connecting to the correct host and correct
## port for the host I'm trying to actually reach (I have 5-6 hosts behind firewalls, using port
## mappings, and this prevents me from having to hardcode those port-mappings in my ~/.ssh/config,
## since the firewall will randomly change those mappings around -- but it at least updates the
## SRV DNS records when it does.  Secondly, it uses some of the nifty features of iTerm2 for Mac
## that lets me use escape codes to set the color of the window tabs and the background color
## of the window to help me find things quicker (when I've got 40 windows open across three
## displays, this makes a huge difference).
##

## Disclaimer: I suck at bash. I should have written this in Perl.  Use at your own risk, or at least
## understand completely what the script does before using it.

## You could use a shell alias to call it, but I kept having weird problems with zsh and autocompletion
## when I did that, so I ended up renaming this script 'ssh' and putting it in my personal ~/bin
## directory, which conveniently is in my PATH before /usr/bin/.  Yes, that's kind of dangerous, so
## you probably shouldn't do that. ;)

## so here we go. please keep your giggling to yourself...

# I'm only pre-initializing these variables because I may append stuff to them later
# and don't need bash getting all uppity if they're missing.

ARGS=""
PORT=""
USER=""
HOST=""
CMD=""

MYNAME=$(basename $0)

# if we're called as 'xsh', it means we want something different. see down below for more.
if [[ $MYNAME == "xsh" ]]; then
    SCREEN=1
else
    SCREEN=0
fi

# loop through the arguments. this is faster/easier than getopt, in my opinion
# since i'm only looking for a few specific things

while [[ -n "$@" ]]; do

    case $1 in
        "-o")
            shift
            ARGS="${ARGS} -o $1"
            shift
            ;;
        "-p")
            shift
            PORT=$1
            shift
            ;;
        "-l")
            shift
            USER=$1
            shift
            ;;
        *@*)
            USER=${1%@*}
            HOST=${1#*@}
            shift
            ;;
        *.*)
            [[ -z $HOST ]] && HOST=$1
            shift
            ;;
        */*)
            CMD=$1
            shift
            ;;
        *\ *)
            CMD=$1
            shift
            ;;
        *)
            [[ -n $ARGS ]] && ARGS="${ARGS} "
            ARGS="${ARGS} ${1}"
            shift
            ;;
    esac
done

## now the main purpose of this script, look up the damned DNS SRV record to find out where I should actually
## be connecting to get to SSH on $HOST

# SSH on the Mac used to do this automatically, then people complained and Apple removed the patch, so now
# we have to look it up ourselves; no big deal, it only takes a fraction of a second.

# let's first save the HOST since we may be nuking it; I want to use the original hostname to calculate the
# colors
ORIGHOST=$HOST

# ... except .local hosts won't have SRV DNS records...
if [[ ${HOST} =~ \.local ]]; then
    HOST=${HOST}
    PORT=""
else
    # search domains be damned; besides, I usually type in FQDN's anyway due to work VPN
    # the dot tremendously speeds things up so the resolver doesn't waste time jumping to every DNS server
    # it can find
    if [[ $HOST =~ \.$ ]]; then
        LOOKUP=${HOST}
    else
        LOOKUP=${HOST}.
    fi

    # now some horriffic piping
    SRV_LOOKUP=$( host -s 1 -t SRV _ssh._tcp.${LOOKUP} 2> /dev/null | grep "has SRV record" | sort -n -k 5,6 | cut -d " " -f 7- | head -n1 )
    if [[ -z $SRV_LOOKUP ]]; then
        # no SRV record; let SSH or .ssh/config figure out what port
        SRV=0
        PORT=""
    else
        # and we now have our destination host and destination port
        PORT=${SRV_LOOKUP%\ *}
        HOST=${SRV_LOOKUP#*\ }
    fi
fi

# because of our . shortcut, we sometimes end up with a host ending with a dot; that messes up the matching
# in ~/.ssh/config, so let's strip it off here
if [[ $HOST =~ (.*)\.$ ]]; then
    HOST=${BASH_REMATCH[1]}
fi

# Now the color stuff comes into play, which only works if you're using iTerm2
if [[ $TERM_PROGRAM == "iTerm.app" ]]; then

    # base colors; I went through the 'safe' HTML color chart and found as many bright colors as different as
    # possible from one another.  I'm colorblind, so my results may be horrible for someone else.
    COLORS=( 0000FF 8A2BE2 A52A2A DEB887 5F9EA0 7FFF00 D2691E FF7F50 6495ED DC143C 00008B 008B8B B8860B A9A9A9 A9A9A9 006400 BDB76B 8B008B 556B2F FF8C00 9932CC 8B0000 E9967A 483D8B 2F4F4F 2F4F4F 00CED1 9400D3 FF1493 00BFFF 696969 696969 1E90FF B22222 228B22 FF00FF FFD700 DAA520 808080 808080 008000 ADFF2F FF69B4 CD5C5C 4B0082 7CFC00 00FF00 32CD32 FF00FF 800000 66CDAA 0000CD BA55D3 9370D8 3CB371 7B68EE 00FA9A 48D1CC C71585 191970 000080 808000 6B8E23 FFA500 FF4500 DA70D6 D87093 CD853F DDA0DD 800080 FF0000 BC8F8F 4169E1 8B4513 FA8072 F4A460 2E8B57 A0522D 6A5ACD 708090 00FF7F 4682B4 D2B48C 008080 FF6347 FFFF00 9ACD32 )

    # this list changes a lot, so.. count 'em
    COLOR_COUNT=${#COLORS[@]}

    # now to get some consistency (same color for each host), let's md5 the hostname, convert a fragment of the md5 to decimal...
    # uses name CRC because I was originally using plain CRC32, but I kept getting numbers that were too close to one another
    if [[ -x /sbin/md5 ]]; then
        CRC=$( echo ${ORIGHOST}${USER}${PORT} | md5 )
        CRC=$( printf "%d" 0x${CRC:0:6} )
    fi

    # goofy math to get a number between 1 and $COLOR_COUNT
    RVAL=$(( $CRC % $COLOR_COUNT ))

    # and get that hex color from the array
    THECOLOR=${COLORS[$RVAL]}

    # now split that color apart into decimal 0-255
    RED=$( printf "%d" 0x${THECOLOR:0:2} )
    GRN=$( printf "%d" 0x${THECOLOR:2:2} )
    BLU=$( printf "%d" 0x${THECOLOR:4:2} )

    # now calculate the window background color; I use 10% of the color because it's barely visible,
    # doesn't screw with the normal ANSI/256 colors, but is visible enough to quickly pick out windows for
    # the same destination across three monitors
    PERCENT=10

    # gotta use bc and printf, since bash can't do floats. this is about the point I should have stopped this
    # project and rewritten it in perl... but... here we go.

    REDBG=$( printf "%0.0f" $( echo "scale=2; $RED * ($PERCENT / 100)" | bc -l ) )
    GRNBG=$( printf "%0.0f" $( echo "scale=2; $GRN * ($PERCENT / 100)" | bc -l ) )
    BLUBG=$( printf "%0.0f" $( echo "scale=2; $GRN * ($PERCENT / 100)" | bc -l ) )

    # now we convert it back into hex, since iTerm is a little inconsistent here (needs hex one place, decimal in another)
    printf "\033]Ph%02x%02x%02x\033\\" $REDBG $GRNBG $BLUBG

    # now we set the window frame and/or tab colors
    printf "\033]6;1;bg;red;brightness;%d\a" $RED
    printf "\033]6;1;bg;green;brightness;%d\a" $GRN
    printf "\033]6;1;bg;blue;brightness;%d\a" $BLU

fi

# build up the known SSH arguments
if [[ -n $USER ]]; then
    ARGS="-l $USER ${ARGS}"
fi
if [[ -n $PORT ]]; then
    ARGS="-p $PORT ${ARGS}"
fi

# if we started this script as 'xsh', we're going to launch a persistent screen session on the remote
# host that this particular session can reconnect to (should I get disconnected). Ideally this would
# use mosh, but mosh's scrollback support sucks dog butt and annoys me...  don't forget the '-tt', or
# ssh will go "f-u, you don't have a TTY" and disconnect.

if [[ $SCREEN -eq 1 ]]; then
    CMD="screen -a -d -RR -S $ITERM_SESSION_ID -e^Bb -i -U -T screen-256color $CMD"
    ARGS="$ARGS -tt"
fi

###########################################################################################################

# have to do a little weirdness here and trap exits, since I need the terminal returned to normal if
# I (for some reason) start to hit ^C before the ssh session is created.. so I use a function and a
# trap to clean things up and exit the script from here on out...

suicide() {
    EC=$?
    if [[ $TERM_PROGRAM == "iTerm.app" ]]; then
        # tab colors to default
        printf '\033]6;1;bg;*;default\a'
        # background color to black
        printf "\033]Ph000000\033\\"
    fi

    # I really have no idea what most of the exit codes mean, and I'm not in the mood to go
    # searching through the OpenSSH source to find out.
    case $EC in
        130)
            echo "SSH exited with exit code 130, because you hit ^C. You hurt its feelings. :("
            ;;
        0)
            # no need to tell them anything. SSH is happy.
            ;;
        *)
            echo "SSH exited with exit code ${EC}."
            ;;
    esac
    exit $EC
}
trap suicide exit

###########################################################################################################

# because dropbox likes to mess with my directory permissions, this is necessary. i used to
# do a 'stat' to see if it was really needed, but it turns out it's just faster to blindly apply the
# changes than do a stat then change them...

chmod 700 $HOME/.ssh > /dev/null 2>&1
chmod 600 $HOME/.ssh/* > /dev/null 2>&1

# now let's call the real SSH... with all of my parameters...
/usr/bin/ssh $ARGS $HOST $CMD
EC=$?

# this 'exit' ends up looping back to my 'suicide' function above
exit $EC