#!/bin/bash

function usage() {
    echo \
"SYNOPSYS
    $0 (-S|--ssh-host) SSHHOST (-t|--tunnel) TUNNELSPEC [OPTIONS]
DESCRIPTION
    Spawn TCP tunnels between your computer and a remote SSH server. If the
    connection fail for some reason, the script will try to re-create the
    tunnels automatically every TIMEOUT seconds.
OPTIONS
    -S,--ssh-host SSHHOST
        Set the ssh host. SSHHOST can be either the full hostname as you would
        specify it on the ssh command-line (e.g. [user@]myserver.com or hostname
        if it is set up in your ssh config).

    -t,--tunnel TUNNELSPEC
        Specify the ports to tunnel as a comma-separated list of either single
        ports or port pairs (e.g. -t 443,80:8080 will establish tunnels
        local:443-remote:443  and local:8080-remote:80). Note that you have to
        be a privileged user (i.e. a user with the CAP_NET_BIND_SERVICE
        capability) to be able to bind to ports below 1024.

    -p,--ssh-port SSHPORT /!\\ Not implemented
        Set the ssh remote port if different from the default.

    -H,--host HOST
        Set the tunnel host. HOST can be either the IP or hostname of the
        server to tunnel to, as seen from the ssh host. Defaults to localhost.

    -i,--interface IF /!\\ Not implemented
        Specify the local interface to bind to. Defaults to localhost. Setting
        it to 0.0.0.0 will make the tunnel available to other clients on your
        network (depending on your firewall rules).

    -f
        Fork the script to the background.

    -h,--help
        Print this help and exit.
EXAMPLES
    Todo
AUTHOR
    Clément Zrounba"
}

function parse_args() {
    ARG_POS_PARAMS=""
    while (( "$#" )); do
        case "$1" in
            -f|--fork)
                ARG_FORK=0
                shift
                ;;
            -h|--help)
                ARG_HELP=0
                shift
                ;;
            -H|--host)
                if [ -n "$2" ] && [ ${2:0:1} != "-" ]; then
                    ARG_HOST=$2
                    shift 2;
                else
                    echo "Error: Argument for $1 is missing" >&2
                    exit 1
                fi
                ;;
            -S|--ssh-host)
                if [ -n "$2" ] && [ ${2:0:1} != "-" ]; then
                    ARG_SSHHOST=$2
                    shift 2;
                else
                    echo "Error: Argument for $1 is missing" >&2
                    exit 1
                fi
                ;;
            -S*)
                if [ -n "${1:2}" ]; then
                    ARG_SSHHOST=${1:2}
                    shift 1;
                else
                    echo "Error: Option not understood: $1" >&2
                    exit 1
                fi
                ;;
            --ssh-host=*)
                if [ -n "${1:2}" ]; then
                    ARG_SSHHOST=${1:11}
                    shift 1;
                else
                    echo "Error: Option not understood: $1" >&2
                    exit 1
                fi
                ;;
            -t|--tunnel)
                if [ -n "$2" ] && [ ${2:0:1} != "-" ]; then
                    ARG_TUNNELSPEC=$2
                    shift 2;
                else
                    echo "Error: Argument for $1 is missing" >&2
                    exit 1
                fi
                ;;
            -*)
                echo "Error: Unsupported flag $1" >&2
                exit 1
                ;;
            *) # preserve positional arguments
                ARG_POS_PARAMS="$ARG_POS_PARAMS $1"
                shift
                ;;
        esac
    done
}

function check_low_port_cap()
{
    [ $(id -u) == 0 ] && echo "yes" && return
    # Check capability as well ?

    echo "no"
}

function parse_tunnel_spec() {
    local tunnel_spec=(${1//,/ })
    local can_access_low_local_ports=$(check_low_port_cap)

    declare -ag _TUNNEL_PORTS_LOCAL
    declare -ag _TUNNEL_PORTS_REMOTE

    for tunnel_spec_entry in ${tunnel_spec[*]}; do
        local ports=(${tunnel_spec_entry/:/ })

        local local_port=${ports[0]}
        local remote_port=${ports[1]:-$local_port}

        echo "Local endpoint $BIND_IF:$local_port will be tunnelled to remote endpoint $TUNNEL_HOST:$remote_port"

        if [ $local_port -lt 1024 ] && [ $can_access_low_local_ports == "no" ]; then
            echo "Error: local port below 1024 specified but running as an unprivileged user" >&2
            exit 1
        fi

        _TUNNEL_PORTS_LOCAL+=($local_port)
        _TUNNEL_PORTS_REMOTE+=($remote_port)
    done
}

# Parse command line arguments into ARG_* variables
parse_args $@

# Make positional parameters accessible through $1, $2, ...
eval set -- "$ARG_POS_PARAMS"

# Print help and exit if requested
if [ 0 == "$ARG_HELP" ]; then
    usage
    exit 0
fi

# Check SSHHOST is given
if [ -z "$ARG_SSHHOST" ]; then
    echo "Error: no ssh host specified" >&2
    exit 1
fi

# Check TUNNELSPEC is given
if [ -z "$ARG_TUNNELSPEC" ]; then
    echo "Error: no tunnel(s) specified" >&2
    exit 1
fi

# Assign values or defaults to other parameters from ARG_* variables
SSH_HOST=${ARG_SSHHOST:-}
BIND_IF=${ARG_BIND_IF:-localhost}
TUNNEL_HOST=${ARG_HOST:-localhost}
SSH_PORT=${ARG_SSHPORT:-22}
DO_FORK=$([ "$ARG_FORK" == "0" ] && echo "yes" || echo "no")

# Parse tunnel args into _TUNNEL_PORTS_* variables
parse_tunnel_spec $ARG_TUNNELSPEC
TUNNEL_PORTS_LOCAL=(${_TUNNEL_PORTS_LOCAL[*]})
TUNNEL_PORTS_REMOTE=(${_TUNNEL_PORTS_REMOTE[*]})

# Logfile location
# comment to use stdout/stderr
#LOGFILE=proxy.log
#LOGFILE_ERR=proxy.log

# set default log to stdout/stderr
LOGFILE=${LOGFILE:-/dev/stdout}
LOGFILE_ERR=${LOGFILE_ERR:-/dev/stderr}

function start_and_monitor_tunnels() {
    # This function takes port numbers as arguments and creates as many TCP
    # tunnels like so:
    #   BIND_IF:$local_port <--> SSH_HOST <--> TUNNEL_HOST:$remote_port
    # If the tunnel fails or ssh exits, we either:
    #   - try again after a timeout if exit code is not 0
    #   - exit if exit code is 0
    # NOTE: for now we just **always** try again after timeout
    local timeout_sec=${1:-10}

    local ntunnels=${#TUNNEL_PORTS_LOCAL[*]}

    declare -a ssh_tunnel_opt
    for i in `seq 0 $(($ntunnels - 1))`; do
        local local_port=${TUNNEL_PORTS_LOCAL[$i]}
        local remote_port=${TUNNEL_PORTS_REMOTE[$i]}

        ssh_tunnel_opt+=("-L")
        ssh_tunnel_opt+=("${BIND_IF}:${local_port}:${TUNNEL_HOST}:${remote_port}")
    done

    while true; do
        # Try establishing the tunnel
        echo -n "Starting tunnel... "
        ssh -N \
            ${ssh_tunnel_opt[*]} \
            -o "ServerAliveInterval 10" \
            -o "ServerAliveCountMax 3" \
            -o "ExitOnForwardFailure yes" \
            ${SSH_HOST} &
        local pid="$!"
        # make sure to kill ssh when killed/parent exits
        trap "kill '$pid'" EXIT
        echo "OK"
        wait

        local exit_code="$?"
        echo "ssh exited with exit code ${exit_code}"
        #if [[ ${exit_code} != 0 ]]; then
        if true; then
            # If ssh exited with non-zero code, wait a bit and retry
            echo "restarting in ${timeout_sec}s..."
            sleep ${timeout_sec}
        else
            echo "exiting."
            return 0
        fi
    done
}

# Start the tunnels and redirect outputs
if [ ${LOGFILE} != "/dev/stdout" ] || [ ${LOGFILE_ERR} != "/dev/stderr" ]; then
    echo "\
    Logs will be redirected to:
        stdout --> ${LOGFILE}
        stderr --> ${LOGFILE_ERR}
    "
fi

# In the following block, the following redirections are set up:
#   - 777 goes to stdout of current script (and not global /dev/stdout)
#   - 1 goes to $LOGFILE
#   - 2 goes to $LOGFILE_ERR
{
    # Enable trace
    #set -x

    # Set up CTRL-C callback (write to stdout of script, not log)
    trap_ctrl_c_callback="
        echo 'Interrupt signal received, exiting...' 1>&777
        exit 130
    "
    trap "${trap_ctrl_c_callback}" SIGINT

    start_and_monitor_tunnels
} 777>&1 1>>${LOGFILE} 2>>${LOGFILE_ERR}