Visual reminder for SSH'ing to production/staging etc servers

If you are anything like me, you have access to a ton of different servers and usually 10 sessions open at once.

If you are anything like me, you have experienced executing a command meant for a staging server on a production machine when under pressure.

The only problem was that when I made that error, that specific webserver was being used to give a presentation to the management and it broke the site...

Hostnames

Usually the hostname in the tab gives you an indication of what type of server it is, but that is not the case for all customers I work for. The customer I was referring to earlier uses a code for the network location and an "auto-increment integer". 

Colors?

So an idea started taking shape: What if I could indicate the type of environment I'm working in with some colors? Colorise the terminal tab, or the background?

Since I'm on MacOS using iTerm2, I looked around to see if I could figure out some kind of visual feedback that would allow me to distinguish the type of server I am on. 

First I discovered that you can modify iTerms' tab colors fairly easy:

tab-color() {
  echo -ne "\033]6;1;bg;red;brightness;$1\a"
  echo -ne "\033]6;1;bg;green;brightness;$2\a"
  echo -ne "\033]6;1;bg;blue;brightness;$3\a"
}

This little gem allows you to specify any color with 3 integer values (RGB) and it will modify the current tab color for you.

You can use it like:

tab-color 255 60 83

To produce some kind of red-ish tab.

Auto triggering

So I made a couple functions for faster access:

tab-red()    { tab-color 255 60 83 }
tab-green()  { tab-color 57 197 77 }
tab-blue()   { tab-color 0 0 255 }
tab-orange() { tab-color 227 143 10 }
tab-reset()  { echo -ne "\033]6;1;bg;*;default\a" }

So now I can just call these to change to the color I want.

But how to get this to automatically trigger?

Using ZSH I looked into the available hooks and it turns out ZSH has some hooks you can use:

precmd
   Executed before each prompt. Note that precommand functions are not re-executed simply because the command line is redrawn, as happens, for example, when a notification about an exiting job is displayed.

preexec
   Executed just after a command has been read and is about to be executed. If the history mechanism is active (regardless of whether the line was discarded from the history buffer), the string that the user typed is passed as the first argument, otherwise it is an empty string. The actual command that will be executed (including expanded aliases) is passed in two different forms: the second argument is a single-line, size-limited version of the command (with things like function bodies elided); the third argument contains the full text that is being executed.

Perfect! precmd can be used to reset the looks and  preexec to evaluate if and what we need to do:

function iterm2_tab_precmd() {
  tab-reset
}

function iterm2_tab_preexec() {
  # ...
}

autoload -U add-zsh-hook
add-zsh-hook precmd  iterm2_tab_precmd
add-zsh-hook preexec iterm2_tab_preexec

Getting things together

The easiest way I could think of to manage which servers belong to what is to just make text-files with lists...

Couple of issues... I usually alias my ssh hosts, so I wanted support for both those aliases and the actual server connected to. Secondly I want it to work for partial matches as well.

An example of one of the files:

sl00012v
dl00132v
staging
acc
pl00403v
sl01401v

So I ended up with this in my .zshrc:

# Only when iTerm is active
if [[ -n "$ITERM_SESSION_ID" ]]; then

  tab-color() {
    echo -ne "\033]6;1;bg;red;brightness;$1\a"
    echo -ne "\033]6;1;bg;green;brightness;$2\a"
    echo -ne "\033]6;1;bg;blue;brightness;$3\a" 
  }

  tab-red()    { tab-color 255 60 83 }
  tab-green()  { tab-color 57 197 77 }
  tab-blue()   { tab-color 0 0 255 }
  tab-orange() { tab-color 227 143 10 }
  tab-reset()  { echo -ne "\033]6;1;bg;*;default\a" }

  # Load both files and just compact the lines into a single '|'-separated string
  SERVERLIST_PRODUCTION=`awk '{print $1}' < ~/.ssh/serverlist.production | paste -s -d\| -`
  SERVERLIST_STAGING=`awk '{print $1}' < ~/.ssh/serverlist.staging | paste -s -d\| -`

  function iterm2_tab_precmd() {
    tab-reset
  }

  function iterm2_tab_preexec() {
    if [[ "$1" =~ "^ssh " ]]; then
      # Assume the host is the last item of the command
      SSH_HOST=`echo $1 | awk '{print $NF}'`
      # Obtain the configuration for connecting to this host and extract the resulting hostname (MacOS only)
      SSH_HOST=`ssh -G $SSH_HOST | grep "^hostname " | awk '{print $2}'`
      # Compare with both the command and the resulting hostname
      if [[ "$1$SSH_HOST" =~ $SERVERLIST_PRODUCTION ]]; then
        tab-red
      elif [[ "$1$SSH_HOST" =~ $SERVERLIST_STAGING ]]; then
        tab-orange
      else
        # Mark SSH session 
        tab-blue 
      fi
    elif [[ "$1" =~ "^./adb " ]]; then
      # Anything in the android debugger is production
      tab-red
    # using my alias for starting a shell in a docker container
    elif [[ "$1" =~ "^docker-shell" ]]; then
      # Separate color for docker containers
      tab-green
    else
      tab-reset
    fi
  }

  autoload -U add-zsh-hook
  add-zsh-hook precmd  iterm2_tab_precmd
  add-zsh-hook preexec iterm2_tab_preexec
fi

Bonus points: Colored background, ... 

While looking I discovered you can also change the iTerm profile for the current tab.
I decided to create profiles for production, staging, external and docker. This also included the colorised tab. The code changes weren't so extensive. Basically all tab-.... function are replaced with the following:

# Change iterm2 profile. Usage iterm2-profile ProfileName (case sensitive)
iterm2-profile() { echo -e "\033]50;SetProfile=$1\a" }
iterm2-default-profile() { iterm2-profile default }

Instead of calling the tab-functions, just call these (eg iterm2-profile production) and you get all changes you want (see the screenshot at the start of the post for an example)

This approach allows you to tinker with the looks without having to dive in your .zshrc file.