#! /bin/sh

VERSION='v0.3.0'

usage() {
    cat <<EOF
Usage: rssht [user@]host[:port] [-f port] [-t port] [-n time] [-p [user@]proxy[:port]] [--http] [-d]
       rssht -h | --help
       rssht --version

Arguments:
    host                   Client hostname
    user                   Username on client host (should have limited rights) [default: same as local username]
    port                   Open port on client host [default: 22 for ssh, 80 for http]

Options:
    -f port                Free port on client host which will redirect to the local host [default: 22000]
    -t port                Open ssh port on local host to forward on the client host [default: 22]
    -k key                 SSH private key to use exclusively on the local host [default: ~/.ssh/rssht_id_rsa]
    -n time                Restore connection time seconds after failure [default: 5]
    -p [user@]proxy[:port] Use the proxy to connect to the client [default username: same as local username; default port: 22]
    --http                 Use ssh over http instead of plain ssh. The client port must forward http traffic to an open ssh port
    -d                     Run as daemon
    -h, --help             Show this screen.
    --version              Show version.
EOF
}

check_args() {
    if test $# -lt 2; then
        echo "Missing argument to option $1" >&2
        exit 1
    fi
}

if test -z "$USER"; then
    USER=$(whoami)
fi

user=
host=
port=
key=
from=22000
to=22
time=5
proxy=
proxy_user=
proxy_host=
proxy_port=
http=false
daemon=false

if test $# -eq 0 || test "x$1" = 'x-h' || test "x$1" = 'x--help'; then
    usage
    exit 0
fi

if test "x$1" = 'x--version'; then
    echo "rssht $VERSION"
    exit 0
fi

while test $# -gt 0; do
    case "$1" in
        -f)
            check_args "$@"
            from="$2"
            shift
        ;;
        -t)
            check_args "$@"
            to="$2"
            shift
        ;;
        -k)
            check_args "$@"
            key="$2"
            shift
        ;;
        -n)
            check_args "$@"
            time="$2"
            shift
        ;;
        -p)
            check_args "$@"
            proxy="$2"
            shift
        ;;
        --http)
            http=true
        ;;
        -d)
            daemon=true
        ;;
        *)
            if test -n "$host"; then
                usage >&2
                exit 1
            fi
            user=$(echo "$1" | sed -n 's/\(.*\)@.*/\1/p')
            host=$(echo "$1" | sed -n 's/\(.*@\)\?\([^:]*\)\(:.*\)\?/\2/p')
            port=$(echo "$1" | sed -n 's/[^:]*:\(.*\)/\1/p')
        ;;
    esac
    shift
done

if test -z "$user"; then
    user="$USER"
fi

if test -z "$host"; then
    usage >&2
    exit 1
fi

if test -z "$port"; then
    if $http; then
        port=80
    else
        port=22
    fi
fi

if test -n "$proxy"; then
    proxy_user=$(echo "$proxy" | sed -n 's/\(.*\)@.*/\1/p')
    proxy_host=$(echo "$proxy" | sed -n 's/\(.*@\)\?\([^:]*\)\(:.*\)\?/\2/p')
    proxy_port=$(echo "$proxy" | sed -n 's/[^:]*:\(.*\)/\1/p')

    if test -z "$proxy_user"; then
        proxy_user="$USER"
    fi

    if test -z "$proxy_port"; then
        proxy_port=22
    fi
fi

if test -z "$key"; then
   if test -f ~/.ssh/rssht_id_rsa; then
       key=~/.ssh/rssht_id_rsa
   fi
fi

# will be wiped at reboot
pid="/tmp/.rssht.$USER.$host.$port.pid"
log="/tmp/.rssht.$USER.$host.$port.log"

set_pid() {
    sed -i "${1}c $2
" "$pid"
}

get_pid() {
    sed -n "${1}p" "$pid"
}

is_running() {
    test -L "/proc/$(get_pid $1)/cwd"
}

random_port() {
    echo $(($1 + $(shuf -i 0-99 -n 1)))
}

free_port() {
    ss -tan | \
        awk -v port=$1 'NR > 1 { gsub(/.*:/, "", $4); if (int($4) >= port) { print $4 } }' | \
        sort -un | \
        awk -v port=$1 '$0 == port { port++; next }; { exit }; END { print port }'
}

err() {
    echo "$1" | tee -a "$log" >&2
}

# skip if already running
if test -f "$pid"; then
    if is_running 1; then
        err "rssht is already running"
        exit 0
    elif is_running 2; then
        err "ssh is already running (pid $(get_pid 2)), can't monitor it"
        exit 1
    fi
fi

# fail if pid file is not writable
> "$log"
echo "$$\n0\n0" > "$pid"
if ! test "x$(cat "$pid")" = "x$(echo "$$\n0\n0")"; then
    err "error: $pid is not writable"
    exit 1
fi

if $daemon; then
    nohup $0 "$user@$host:$port" -f "$from" -t "$to" -n "$time" $($http && echo "--http") >/dev/null 2>&1 &
    exit 0
fi

# FIXME The following approach to get a free port is unfortunately not race-free since
# the free port detection is not atomic (hence the randomization, to lower the race
# probability). This is because we can't reserve a free port for future use by SSH.
http_port=$(free_port $(random_port 8022))

monitor="-o ServerAliveInterval=30 -o ServerAliveCountMax=3"

set_pid 1 "$$"
err "[$(date)] starting"
while true; do
    if $http; then
        if ! is_running 3; then
            htc -F $http_port $host:$port >> "$log" 2>&1
            set_pid 3 "$!"
        fi
        cmd_host=localhost
        cmd_port=$http_port
        cmd_alias="-o HostKeyAlias=rssht_${host}_$port"
    else
        cmd_host=$host
        cmd_port=$port
        cmd_alias=
    fi
    if test -n "$proxy"; then
        if test -z "$key"; then
            # TODO FIXME does not block (because no tty?)
            ssh -nxNTC $monitor $cmd_alias -o ProxyCommand="ssh $monitor $cmd_alias $proxy_user@$proxy_host -p $proxy_port -W $cmd_host:$cmd_port" -R$from:localhost:$to $user@$cmd_host -p $cmd_port >> "$log" 2>&1 &
        else
            # TODO FIXME does not block (because no tty?)
            ssh -nxNTC $monitor $cmd_alias -o ProxyCommand="ssh $monitor $cmd_alias -o ControlMaster=no -o IdentitiesOnly=yes -i $key $proxy_user@$proxy_host -p $proxy_port -W $cmd_host:$cmd_port" -o ControlMaster=no -o IdentitiesOnly=yes -i "$key" -R$from:localhost:$to $user@$cmd_host -p $cmd_port >> "$log" 2>&1 &
        fi
    else
        if test -z "$key"; then
            # TODO FIXME does not block (because no tty?)
            ssh -nxNTC $monitor $cmd_alias $cmd_proxy -R$from:localhost:$to $user@$cmd_host -p $cmd_port >> "$log" 2>&1 &
        else
            # TODO FIXME does not block (because no tty?)
            ssh -nxNTC $monitor $cmd_alias $cmd_proxy -o ControlMaster=no -o IdentitiesOnly=yes -i "$key" -R$from:localhost:$to $user@$cmd_host -p $cmd_port >> "$log" 2>&1 &
        fi
    fi
    set_pid 2 "$!"
    wait %1
    err "[$(date)] stopped"
    sleep "$time"
    err "[$(date)] restarting"
done

# TODO trap ctrl+c, otherwise this will never happen
rm "$pid"