#!/bin/bash

CONFDIR="$HOME/.config/ssh-tunnelier"
CONF="$CONFDIR/tunnels.conf"
# Just over a year in minutes
MAGIC_TIME=525601

BASE="base64 -w 0"
UNBASE="base64 -d"
if [[ "$OSTYPE" = "darwin"* ]]; then
  BASE="base64"
  UNBASE="base64 -D"
fi
LOCALHOSTSYMBOL="💻"

function _helpexit() {
  echo "SSH tunnel manager
Runs and monitors preconfigured background ssh tunnels, with the ability
to kill existing ssh processes. You don't need to keep the program on
to remember connected sessions.

Configuration stored in: $CONF

If run with arguments, runs an ssh tunnel with arguments:
  ssh-tunnelier hostname [local-port] remote-port

If local port missing, the same port assumed in local
  "
  exit 0
}

for (( i=1; i<=$#; i++ )); do
  [[ ${!i} = "-h" ]] && _helpexit
  [[ ${!i} = "--help" ]] && _helpexit
done

[[ -f "$CONF" ]]  || {
  echo No config file.
  echo "$CONF" each line like:
  echo "servername: 8888,8889:8989,8890:host:9090"
  mkdir -p "$CONFDIR"
  echo -e "# example:\nservername: 8888,8889:8989,8890:host:9090" > "$CONF"
}
number_re='^[0-9]+$'
# import qolop
_qCol(){ true; } # incase qolop missing
. qolop &>/dev/null
colZ=$( _qCol z )
colTitle=$( _qCol z G )
colRow=$( _qCol z S )
colMenu=$( _qCol z y )
colWarn=$( _qCol z R )

function fpgrep() {
  [[ -z $( pgrep -x -u $UID "$@" ) ]] && { echo '----'; return; }
  pgrep -x -u $UID "$@" | xargs ps -o pid,start,args
}

function get_id() {
  id=$( printf "%s" "$1" | $BASE )
  echo -n $id
}

function get_command() {
  switches=$( echo "$1" | $UNBASE | parse_switches )
  echo -n "ssh -f -n $switches \"nice /bin/bash -c 'for ((i=1;i<$MAGIC_TIME;i++)); do cut -f4 -d \\\" \\\" /proc/\\\$PPID/stat | xargs kill -0 || exit ; sleep 60;done'; echo tunnelier $1\""
}

function parse_switches() {
  python3 -c "
import sys
host, tunnels = sys.stdin.read(1024*10).split(':',1)

for i, tunnel in enumerate(tunnels.split(',')):
    tunnel = tunnel.strip().split(':')
    thost='localhost'
    if len(tunnel) == 1:
        tport=tunnel[0]
    if len(tunnel) == 2:
        tport=tunnel[1]
    if len(tunnel) == 3:
        thost=tunnel[1]
        tport=tunnel[2]
    sys.stdout.write('-L {}:{}:{} '.format(tunnel[0], thost, tport))
    sys.stderr.write('Generated tunnel {}: http://localhost:{}  ->  {}:{}\n'.format(i+1,tunnel[0],thost, tport))
sys.stdout.write(host)
  "
}

function get_pid() {
  pgrep -f "$MAGIC_TIME.*echo tunnelier $1" | head -n 1
}

function list_all_ssh() {
  printf "\n${colTitle}=============================================\n"
  printf "${colTitle}                List of SSH:${colZ}\n"
  printf "${colTitle}Enter PID to kill, empty to return $colZ\n"
  printf "${colTitle}Outgoing SSH processes\n"
  printf "${colTitle}=============================================\n${colRow}"
  fpgrep ssh
  printf "${colTitle}Incoming SSH processes\n"
  printf "${colTitle}=============================================\n${colRow}"
  fpgrep sshd
  read -t 600 inputpid

  [[ "$inputpid" =~ $number_re ]] && {
    ask_to_kill $inputpid
  }
}

function reset_all_ssh() {
  match="ssh.*sleep.*echo tunnelier"
  printf "\n${colTitle}=============================================\n"
  printf "${colTitle}                List of tunneliers:${colZ}\n"
  printf "${colTitle}=============================================\n${colRow}"
  if [[ -z $( pgrep -u $UID -f "$match" ) ]]; then
    echo 'No tunneliers running.'
    return
  else
    pgrep -u $UID -f "$match" | xargs ps -o pid,start,args
  fi
  printf "\n  k kill all\n  t terminate all\n  empty returns${colZ}\n"
  read -t 600 input1
  [[ "$input1" = "k" ]] && pkill -u $UID -f "$match"
  [[ "$input1" = "t" ]] && kill -u $UID -9 -f "$match"
}

function read_config() {
  while read line; do
    id=$( get_id "$line" )
    pid=$( get_pid "$id" )
    printf "%-3d %-7s %s\n" $i "$pid" "$line" | sed 's,:localhost:,:'$LOCALHOSTSYMBOL':,g'
    ids+=( $id )
    i=$(( i + 1 ))
  done < <( grep -v ^# "$CONF" | grep '[a-zA-Z]' )
}

function run_args() {
  if [[ -z "$2" ]];  then
    # only hostname
    if grep -q "^$1:" "$CONF"; then
      grep "^$1:" "$CONF"
      host="$1"
      line=$( grep "^$1:" "$CONF" )
      switches=$( echo "$line" | parse_switches )
    else
      echo No such host
      cat "$CONF"
      exit 1
    fi
  else
    host="$1"
    shift 1
    line="$host: $@"
    switches=$( echo "$line" | parse_switches )
  fi
  id=$( get_id "$line" )
  eval "ssh -f -n $switches \"nice /bin/bash -c 'for ((i=1;i<$MAGIC_TIME;i++)); do cut -f4 -d \\\" \\\" /proc/\\\$PPID/stat | xargs kill -0 || exit ; sleep 60;done'; echo tunnelier $id\""
  exit $?
}

function run_command() {
  eval $( get_command "$1" )
}

function ask_to_kill() {
  printf "${colRow}"
  ps -o pid,bsdstart,args "$1"
  printf "\n  k kill\n  t terminate${colZ}\n"
  read -t 600 input2
  [[ "$input2" = "k" ]] && kill $1
  [[ "$input2" = "t" ]] && kill -9 $1
}

function instant_entry() {
  printf "${colMenu}Start a local forward tunnel${colZ}\n"
  printf "${colMenu}Host name?${colZ}\n"
  read -e HOST
  test -z "$HOST" && return
  printf "${colMenu}Target port at host?${colZ}\n"
  read -e REMOTE
  test -z "$REMOTE" && return
  printf "${colMenu}Local port where tunnel starts?${colZ}\n"
  read -e -i "$REMOTE" LOCAL
  test -z "$LOCAL" && return
  printf "${colRow}-L ${LOCAL}:localhost:${REMOTE} ${HOST}${colZ}\n"
  run_args "$HOST" "$LOCAL:$REMOTE"
  printf "${colMenu}See the connection in 'list ssh'${colZ}\n"
}

if [[ -n "$1" ]]; then
  run_args "$@"
  exit $?
fi

while true; do
  ids=()
  i=1
  printf "\n${colTitle}=============================================\n"
  printf "${colTitle}   List of tunnels: (${LOCALHOSTSYMBOL}=localhost)    [%s]${colZ}\n" $( date +%H:%M )
  printf "${colTitle} (q)uit (e)dit (l)ist (i)nstant (r)eset${colZ}\n"
  printf "${colTitle}=============================================\n"
  printf "${colRow}ID  PID     host: ports\n"
  read_config
  printf "$colZ"
  read -t 600 input

  [[ "$input" = "q" ]] && exit 0
  [[ "$input" = "i" ]] && instant_entry
  [[ "$input" = "e" ]] && vim "$CONF"
  [[ "$input" = "l" ]] && list_all_ssh
  [[ "$input" = "r" ]] && reset_all_ssh
  [[ "$input" = "0" ]] && continue
  [[ "$input" =~ $number_re ]] && {
    j=$(( $input - 1 ))
    [[ $j -gt $(( $i - 1 )) ]] && {
      printf  "${colWarn}No such tunnel number${colZ}\n"
      continue
    }

    printf "\n${colTitle}Tunnel command: %s${colZ}\n" "$( echo  "${ids[$j]}" | $UNBASE )"
    this_pid="$( get_pid "${ids[$j]}" )"
    [[ -n "$this_pid" ]] && {
      # PID exists, ask to kill
      ask_to_kill "$this_pid"
    }
    [[ -z "$this_pid" ]] && {
      # PID empty, run
      run_command "${ids[$j]}"
      sleep 1
    }
  }
done