#!/usr/bin/env bash

REALPATH="$(realpath "$0")"
BASEDIR="${REALPATH%/*}"

SSHCONF="$BASEDIR/var/ssh.conf"
PWCONF="$BASEDIR/var/passwords.conf"
WATCHCONF="$BASEDIR/var/watch.conf"
DEBUGCONF="$BASEDIR/var/debug.conf"

function genconfig () {
  mkdir -p "$BASEDIR/var/log"
  cat <<EOF > "$SSHCONF"
Host *
  ExitOnForwardFailure yes
  TCPKeepAlive yes
  ServerAliveInterval 15
  ServerAliveCountMax 3
  ControlMaster auto
  StrictHostKeyChecking no

EOF
  cat "$BASEDIR/etc/tunnelkeeper.conf" | sed 's/#.*$//g' | grep -v '^$' | grep -Eiv '^ *(watch|password|debug)' >> "$SSHCONF"
  cat "$BASEDIR/etc/tunnelkeeper.conf" | grep -Ei '^( *watch|Host)' | awk '{print $1 " " $2}' | grep -i -B1 --no-group-separator watch | tr '\n' ' ' | sed 's/Host /\n/g; s/ *[Ww]atch//g' > "$WATCHCONF"
  cat "$BASEDIR/etc/tunnelkeeper.conf" | grep -Ei '^( *password|Host)' | awk '{print $1 " " $2}' | grep -i -B1 --no-group-separator password | tr '\n' ' ' | sed 's/Host /\n/g; s/ *[Pp]assword//g' > "$PWCONF"
  cat "$BASEDIR/etc/tunnelkeeper.conf" | grep -Ei '^( *debug|Host)' | awk '{print $1 " " $2}' | grep -i -B1 --no-group-separator debug | tr '\n' ' ' | sed 's/Host /\n/g; s/ *[Dd]ebug//g' > "$DEBUGCONF"
}

function ruroot () {
  if [[ $UID -ne 0 ]]; then
    echo "You must be root to do this"
    exit
  fi
}

function debugopt () {
  case "$(awk "/^$1/ {print \$2}" ${DEBUGCONF})" in
    2) echo -n '-v';;
    3) echo -n '-vvv';;
    *) echo -n '';;
  esac
}

function dbg () {
  if [[ "$(awk "/^$1/ {print \$2}" ${DEBUGCONF})" =~ [123] ]]; then
    cat | sed "s/^/$(date +"%H:%M:%S")  /g" >> "$BASEDIR/var/log/tunnelkeeper-$(date +"%Y%m%d").log"
  fi
}

function connect () {
  bo="$(debugopt $1)"
  ssh -F "${SSHCONF}" $bo -S "$BASEDIR/var/$1.tksock" -N $1 '#tunnelkeeper' 2>&1 | dbg $1
}

function forkwatch () {
  host=$1
  while true; do
    TIMEOUT=$(awk "/^$host/ {print \$2}" ${WATCHCONF})
    sleep $TIMEOUT
    echo "tick $host" | dbg
    if [[ $(timeout $TIMEOUT ssh localhost -o "StrictHostKeyChecking no" -S $BASEDIR/var/$host.tksock "echo tk") != "tk" ]]; then
      ssh localhost -O exit -S $BASEDIR/var/$host.tksock
      echo "Killing connection to $host. Trying again." | dbg $host
    fi
  done
}

function forkstart () {
  host=$1
  (grep "$host" $WATCHCONF &>/dev/null) && forkwatch $host &
  if [[ $(grep -c "^$host" $PWCONF) -gt 0 ]]; then
    pass=$(awk "/^$host/ {print \$2}" ${PWCONF}) # password needed
    screen -d -m -S "tk${host}" $0 FORKSCREEN $host '#tunnelkeeper'
    while true; do
      sleep 5
      if [[ -f "$BASEDIR/var/${host}.screen" ]]; then
        screen -S "tk${host}" -X stuff "$pass
"
        rm -f "$BASEDIR/var/${host}.screen"
      fi
    done
  else # passwordless auth
    while true; do
      connect $host
      sleep 5
    done
  fi
}

### main loop

case "$1" in
#   FORKSTART )
#     (grep "$2" $WATCHCONF &>/dev/null) && $0 FORKWATCH $2 '#tunnelkeeper' &
#     if [[ $(grep -c "^$2" $PWCONF) -gt 0 ]]; then
#       pass=$(awk "/^$2/ {print \$2}" ${PWCONF}) # password needed
#       screen -d -m -S "tk${2}" $0 FORKSCREEN $2 '#tunnelkeeper'
#       while true; do
#         sleep 5
#         if [[ -f "$BASEDIR/var/${2}.screen" ]]; then
#           screen -S "tk${2}" -X stuff "$pass
# "
#           rm -f "$BASEDIR/var/${2}.screen"
#         fi
#       done
#     else # passwordless auth
#       while true; do
#         connect $2
#         sleep 5
#       done
#     fi
#     exit
#   ;;
  # FORKWATCH ) # makes sure the connection is still working, even if ssh doesn't drop it. Needs a login shell to work.
  #   while true; do
  #     TIMEOUT=$(awk "/^$2/ {print \$2}" ${WATCHCONF})
  #     sleep $TIMEOUT
  #     echo "tick $2" | dbg
  #     if [[ $(timeout $TIMEOUT ssh localhost -o "StrictHostKeyChecking no" -S $BASEDIR/var/$2.tksock "echo tk") != "tk" ]]; then
  #       ssh localhost -O exit -S $BASEDIR/var/$2.tksock
  #       echo "Killing connection to $2. Trying again." | dbg $2
  #     fi
  #   done
  #   exit
  # ;;
  FORKSCREEN )
    while true; do
      echo $$ > "$BASEDIR/var/${2}.screen"
      connect $2 #$dbgopt
      sleep 5
    done
    exit # de-forkbombing exit.
  ;;
  start)
    [[ -e "$BASEDIR/var/tunnelkeeper.pid" ]] && exit
    genconfig
    echo $$ > "$BASEDIR/var/tunnelkeeper.pid"
    # cat "$BASEDIR/etc/tunnelkeeper.conf" | awk '/^Host / {print $2}' | xargs -I% -P0 $0 FORKSTART % '#tunnelkeeper' &
    cat "$BASEDIR/etc/tunnelkeeper.conf" | awk '/^Host / {print $2}' | while read host
      do forkstart $host &
    done
  ;;
  stop )
    rm -f "$BASEDIR/var/tunnelkeeper.pid"
    rm -f "$BASEDIR/var/*.conf"
    pkill -f '#tunnelkeeper' &>/dev/null
  ;;
  restart ) 
    if systemctl status tunnelkeeper &>/dev/null; then
      systemctl restart tunnelkeeper
    else
      $0 stop; sleep 2; $0 start
    fi
  ;;
  install )
    ruroot
    which screen &>/dev/null || yum install -y screen || apt install -y screen || echo "Couldn't install screen"
    mkdir -p /opt/tunnelkeeper/var
    mkdir -p /opt/tunnelkeeper/etc
    if [[ "$BASEDIR" != '/opt/tunnelkeeper/' ]]; then
      cp -n "$BASEDIR/etc/tunnelkeeper.conf" /opt/tunnelkeeper/etc/
      cp "$REALPATH" "/opt/tunnelkeeper/tunnelkeeper"
    fi
    ln -f -s /opt/tunnelkeeper/tunnelkeeper /usr/local/bin/tunnelkeeper
    echo "[Unit]
Description=TunnelKeeper SSH tunnel utility.
After=network.target
[Service]
User=root
Group=root
Type=forking
ExecStart=/opt/tunnelkeeper/tunnelkeeper start
ExecStop=/opt/tunnelkeeper/tunnelkeeper stop
RestartSec=15
Restart=always
[Install]
WantedBy=multi-user.target" > /lib/systemd/system/tunnelkeeper.service
    systemctl daemon-reload
    systemctl enable tunnelkeeper.service
    echo -e "\nTunnelKeeper service installed.\n"
  ;;
  uninstall )
    ruroot
    systemctl disable tunnelkeeper.service
    rm /usr/local/bin/tunnelkeeper
    rm /lib/systemd/system/tunnelkeeper.service
    systemctl daemon-reload
    cp -f /opt/tunnelkeeper/etc/tunnelkeeper.conf /opt/tunnelkeeper.conf.bak
    rm -rf /opt/tunnelkeeper
    echo -e "\nTunnelKeeper service uninstalled. Config backup saved as /opt/tunnelkeeper.conf.bak \n"
  ;;
  list )
    echo "---"
    find "$BASEDIR/var/" -name '*.tksock' | sed 's/^.*\///g; s/\.tksock//g'
    echo "---"
  ;;
  config )
    ruroot
    vi "$BASEDIR/etc/tunnelkeeper.conf"
  ;;
   * )
    echo -e "\nUsage: $(basename $0) start|stop|restart|install|uninstall|config|list\n"
  ;;
esac