#!/bin/ksh

# Copyright (C) 2018 The noVNC authors
# Licensed under MPL 2.0 or any later version (see LICENSE.txt)

usage() {
    if [ "$*" ]; then
        echo "$*"
        echo
    fi
    echo "Usage: ${NAME} [--listen [HOST:]PORT] [--vnc VNC_HOST:PORT] [--cert CERT] [--ssl-only]"
    echo
    echo "Starts the WebSockets proxy and a mini-webserver and "
    echo "provides a cut-and-paste URL to go to."
    echo
    echo "    --listen [HOST:]PORT  Port for proxy/webserver to listen on"
    echo "                          Default: 6080 (on all interfaces)"
    echo "    --vnc VNC_HOST:PORT   VNC server host:port proxy target"
    echo "                          Default: localhost:5900"
    echo "    --cert CERT           Path to combined cert/key file, or just"
    echo "                          the cert file if used with --key"
    echo "                          Default: self.pem"
    echo "    --key KEY             Path to key file, when not combined with cert"
    echo "    --web WEB             Path to web files (e.g. vnc.html)"
    echo "                          Default: ./"
    echo "    --ssl-only            Disable non-https connections."
    echo "                                    "
    echo "    --file-only           Disable directory listing in web server."
    echo "                                    "
    echo "    --record FILE         Record traffic to FILE.session.js"
    echo "                                    "
    echo "    --syslog SERVER       Can be local socket such as /dev/log, or a UDP host:port pair."
    echo "                                    "
    echo "    --heartbeat SEC       send a ping to the client every SEC seconds"
    echo "    --timeout SEC         after SEC seconds exit when not connected"
    echo "    --idle-timeout SEC    server exits after SEC seconds if there are no"
    echo "                                    "
    echo "    --web-auth            enable authentication"
    echo "    --auth-plugin CLASS   authentication plugin to use"
    echo "    --auth-source ARG     plugin configuration"
    echo "                                    "
    echo "                          active connections"
    echo "                                    "
    exit 2
}

NAME="$(basename $0)"
REAL_NAME="$(readlink -f $0)"
HERE="$(cd "$(dirname "$REAL_NAME")" && pwd)"
HOST=""
PORT="6080"
LISTEN="$PORT"
VNC_DEST="localhost:5900"
CERT=""
KEY=""
WEB=""
proxy_pid=""
SSLONLY=""
RECORD=""
SYSLOG_ARG=""
HEARTBEAT_ARG=""
IDLETIMEOUT_ARG=""
TIMEOUT_ARG=""
WEBAUTH_ARG=""
AUTHPLUGIN_ARG=""
AUTHSOURCE_ARG=""
FILEONLY_ARG=""


die() {
    echo "$*"
    exit 1
}

cleanup() {
    trap - TERM QUIT INT EXIT
    trap "true" CHLD   # Ignore cleanup messages
    echo
    if [ -n "${proxy_pid}" ]; then
        echo "Terminating WebSockets proxy (${proxy_pid})"
        kill ${proxy_pid}
    fi
}

# Process arguments

# Arguments that only apply to chrooter itself
while [ "$*" ]; do
    param=$1; shift; OPTARG=$1
    case $param in
    --listen)  LISTEN="${OPTARG}"; shift            ;;
    --vnc)     VNC_DEST="${OPTARG}"; shift        ;;
    --cert)    CERT="${OPTARG}"; shift            ;;
    --key)     KEY="${OPTARG}"; shift             ;;
    --web)     WEB="${OPTARG}"; shift            ;;
    --ssl-only) SSLONLY="--ssl-only"             ;;
    --file-only) FILEONLY_ARG="--file-only"      ;;
    --record) RECORD="${OPTARG}"; shift ;;
    --syslog) SYSLOG_ARG="--syslog ${OPTARG}"; shift ;;
    --heartbeat) HEARTBEAT_ARG="--heartbeat ${OPTARG}"; shift ;;
    --idle-timeout) IDLETIMEOUT_ARG="--idle-timeout ${OPTARG}"; shift ;;
    --timeout) TIMEOUT_ARG="--timeout ${OPTARG}"; shift ;;
    --web-auth) WEBAUTH_ARG="--web-auth"                ;;
    --auth-plugin) AUTHPLUGIN_ARG="--auth-plugin ${OPTARG}"; shift ;;
    --auth-source) AUTHSOURCE_ARG="--auth-source ${OPTARG}"; shift ;;
    -h|--help) usage                              ;;
    -*) usage "Unknown chrooter option: ${param}" ;;
    *) break                                      ;;
    esac
done

if [ "$LISTEN" != "$PORT" ]; then
    HOST=${LISTEN%:*}
    PORT=${LISTEN##*:}
    # if no host was given, restore
    [ "$HOST" = "$PORT" ] && HOST=""
fi

# Sanity checks
if [ -z "${HOST}" ]; then
    netstat -ln -p tcp | grep -qs "\.${PORT} .*LISTEN" \
        && die "Port ${PORT} in use. Try --listen PORT"
fi

trap "cleanup" TERM QUIT INT EXIT

# Find vnc.html
if [ -n "${WEB}" ]; then
    if [ ! -e "${WEB}/vnc.html" ]; then
        die "Could not find ${WEB}/vnc.html"
    fi
elif [ -e "$(pwd)/vnc.html" ]; then
    WEB=$(pwd)
elif [ -e "${HERE}/../vnc.html" ]; then
    WEB=${HERE}/../
elif [ -e "${HERE}/vnc.html" ]; then
    WEB=${HERE}
elif [ -e "${HERE}/../share/novnc/vnc.html" ]; then
    WEB=${HERE}/../share/novnc/
else
    die "Could not find vnc.html"
fi

# Find self.pem
if [ -n "${CERT}" ]; then
    if [ ! -e "${CERT}" ]; then
        die "Could not find ${CERT}"
    fi
elif [ -e "$(pwd)/self.pem" ]; then
    CERT="$(pwd)/self.pem"
elif [ -e "${HERE}/../self.pem" ]; then
    CERT="${HERE}/../self.pem"
elif [ -e "${HERE}/self.pem" ]; then
    CERT="${HERE}/self.pem"
else
    echo "Warning: could not find self.pem"
fi

# Check key file
if [ -n "${KEY}" ]; then
    if [ ! -e "${KEY}" ]; then
        die "Could not find ${KEY}"
    fi
fi

WEBSOCKIFY=/usr/local/bin/websockify

# Make all file paths absolute as websockify changes working directory
WEB=`realpath "${WEB}"`
[ -n "${CERT}" ] && CERT=`realpath "${CERT}"`
[ -n "${KEY}" ] && KEY=`realpath "${KEY}"`
[ -n "${RECORD}" ] && RECORD=`realpath "${RECORD}"`

echo "Starting webserver and WebSockets proxy on${HOST:+ host ${HOST}} port ${PORT}"
${WEBSOCKIFY} ${SYSLOG_ARG} ${SSLONLY} ${FILEONLY_ARG} --web ${WEB} ${CERT:+--cert ${CERT}} ${KEY:+--key ${KEY}} ${LISTEN} ${VNC_DEST} ${HEARTBEAT_ARG} ${IDLETIMEOUT_ARG} ${RECORD:+--record ${RECORD}} ${TIMEOUT_ARG} ${WEBAUTH_ARG} ${AUTHPLUGIN_ARG} ${AUTHSOURCE_ARG} &
proxy_pid="$!"
sleep 1
if [ -z "$proxy_pid" ] || ! ps -eo pid= | grep -w "$proxy_pid" > /dev/null; then
    proxy_pid=
    echo "Failed to start WebSockets proxy"
    exit 1
fi

if [ -z "$HOST" ]; then
    HOST=$(hostname)
fi

echo -e "\n\nNavigate to this URL:\n"
if [ "x$SSLONLY" == "x" ]; then
    echo -e "    http://${HOST}:${PORT}/vnc.html?host=${HOST}&port=${PORT}\n"
else
    echo -e "    https://${HOST}:${PORT}/vnc.html?host=${HOST}&port=${PORT}\n"
fi

echo -e "Press Ctrl-C to exit\n\n"

wait ${proxy_pid}
