#!/bin/bash
### every exit != 0 fails the script
set -e
DEBUG=true
no_proxy="localhost,127.0.0.1"

# dict to store processes
declare -A KASM_PROCS

# switch passwords to local variables
tmpval=$VNC_VIEW_ONLY_PW
unset VNC_VIEW_ONLY_PW
VNC_VIEW_ONLY_PW=$tmpval
tmpval=$VNC_PW
unset VNC_PW
VNC_PW=$tmpval

STARTUP_COMPLETE=0

######## FUNCTION DECLARATIONS ##########

## print out help
function help (){
    echo "
        USAGE:

        OPTIONS:
        -w, --wait      (default) keeps the UI and the vncserver up until SIGINT or SIGTERM will received
        -s, --skip      skip the vnc startup and just execute the assigned command.
                        example: docker run kasmweb/core --skip bash
        -d, --debug     enables more detailed startup output
                        e.g. 'docker run kasmweb/core --debug bash'
        -h, --help      print out this help

        Fore more information see: https://github.com/ConSol/docker-headless-vnc-container
        "
}

## correct forwarding of shutdown signal
function cleanup () {
    kill -s SIGTERM $!
    exit 0
}

function start_kasmvnc (){
    if [[ $DEBUG == true ]]; then
      echo -e "\n------------------ Start KasmVNC Server ------------------------"
    fi

    if [ "x$BASE_PORT" == "x" ]; then
        BASE_PORT=8590
        DESKTOP_NUMBER=1
    else
        DESKTOP_NUMBER=$(($BASE_PORT-5900+1))
    fi
    
    #if [ "x$KASMSOCK" == "xTrue" ]; then
    #    export SOCKET_PORT=$(( $RANDOM % 50 + 1 ))
    #fi

    echo "HERE"

    /usr/local/bin/vncserver :$DESKTOP_NUMBER -depth 24 -geometry 1280x1050 -websocketPort $BASE_PORT -cert ${HOME}/.vnc/self.pem -httpd /usr/local/share/kasmvnc/www -disableBasicAuth -FrameRate=24 -interface 0.0.0.0
    echo "HEREDONE"

    KASM_PROCS['kasmvnc']=$(cat $HOME/.vnc/*${DISPLAY_NUM}.pid)

    if [[ $DEBUG == true ]]; then
      echo -e "\n------------------ Started Websockify  ----------------------------"
      echo "Websockify PID: ${KASM_PROCS['kasmvnc']}";
    fi
}

function start_window_manager (){
    echo -e "start window manager\n..."
    $STARTUPDIR/window_manager_startup.sh #&> $STARTUPDIR/window_manager_startup.log
}

function start_audio_out_websocket (){
    if [[ ${KASM_SVC_AUDIO:-1} == 1 ]]; then
        echo 'Starting audio websocket server'
        $STARTUPDIR/jsmpeg/kasm_audio_out-linux kasmaudio 8081 4901 ${HOME}/.vnc/self.pem ${HOME}/.vnc/self.pem "kasm_user:$VNC_PW"  &

        KASM_PROCS['kasm_audio_out_websocket']=$!

        if [[ $DEBUG == true ]]; then
          echo -e "\n------------------ Started Audio Out Websocket  ----------------------------"
          echo "Kasm Audio Out Websocket PID: ${KASM_PROCS['kasm_audio_out_websocket']}";
        fi
    fi
}

function start_audio_out (){
    if [[ ${KASM_SVC_AUDIO:-1} == 1 ]]; then
        echo 'Starting audio server'

        if [ "${START_PULSEAUDIO:-0}" == "1" ] ;
        then
            echo "Starting Pulse"
            pulseaudio --start
        fi

        if [[ $DEBUG == true ]]; then
            echo 'Starting audio service in debug mode'
            no_proxy=127.0.0.1 ffmpeg -f pulse -fragment_size ${PULSEAUDIO_FRAGMENT_SIZE:-2000} -ar 44100 -i default -f mpegts -correct_ts_overflow 0 -codec:a mp2 -b:a 128k -ac 1 -muxdelay 0.001 http://127.0.0.1:8081/kasmaudio &
            KASM_PROCS['kasm_audio_out']=$!
        else
            echo 'Starting audio service'
            no_proxy=127.0.0.1 ffmpeg -v verbose -f pulse -fragment_size ${PULSEAUDIO_FRAGMENT_SIZE:-2000} -ar 44100 -i default -f mpegts -correct_ts_overflow 0 -codec:a mp2 -b:a 128k -ac 1 -muxdelay 0.001 http://127.0.0.1:8081/kasmaudio > /dev/null 2>&1 &
            KASM_PROCS['kasm_audio_out']=$!
            echo -e "\n------------------ Started Audio Out  ----------------------------"
            echo "Kasm Audio Out PID: ${KASM_PROCS['kasm_audio_out']}";
        fi
    fi
}

function start_audio_in (){
    if [[ ${KASM_SVC_AUDIO_INPUT:-1} == 1 ]]; then
        echo 'Starting audio input server'
        $STARTUPDIR/audio_input/kasm_audio_input_server --ssl --auth-token "kasm_user:$VNC_PW" --cert ${HOME}/.vnc/self.pem --certkey ${HOME}/.vnc/self.pem &

        KASM_PROCS['kasm_audio_in']=$!

        if [[ $DEBUG == true ]]; then
            echo -e "\n------------------ Started Audio Out Websocket  ----------------------------"
            echo "Kasm Audio In PID: ${KASM_PROCS['kasm_audio_in']}";
        fi
    fi
}

function start_upload (){
    if [[ ${KASM_SVC_UPLOADS:-1} == 1 ]]; then
        echo 'Starting upload server'
        cd $STARTUPDIR/upload_server/
        ./kasm_upload_server --ssl --auth-token "kasm_user:$VNC_PW" &

        KASM_PROCS['upload_server']=$!

        if [[ $DEBUG == true ]]; then
            echo -e "\n------------------ Started Audio Out Websocket  ----------------------------"
            echo "Kasm Audio In PID: ${KASM_PROCS['upload_server']}";
        fi
    fi
}

############ END FUNCTION DECLARATIONS ###########

if [[ $1 =~ -h|--help ]]; then
    help
    exit 0
fi

# should also source $STARTUPDIR/generate_container_user
source $HOME/.bashrc

if [[ ${KASM_DEBUG:-0} == 1 ]]; then
    echo -e "\n\n------------------ DEBUG KASM STARTUP -----------------"
    export DEBUG=true
    set -x
fi

trap cleanup SIGINT SIGTERM

## resolve_vnc_connection
VNC_IP=$(hostname -i)
if [[ $DEBUG == true ]]; then
  echo "IP Address used for external bind: $VNC_IP"
fi

# Create cert for KasmVNC
#mkdir -p ${HOME}/.vnc
#openssl req -x509 -nodes -days 3650 -newkey rsa:2048 -keyout ${HOME}/.vnc/self.pem -out ${HOME}/.vnc/self.pem -subj "/C=US/ST=VA/L=None/O=None/OU=DoFu/CN=kasm/emailAddress=none@none.none"

# first entry is control, second is view (if only one is valid for both)
mkdir -p "$HOME/.vnc"
PASSWD_PATH="$HOME/.kasmpasswd"
if [[ -f $PASSWD_PATH ]]; then
    echo -e "\n---------  purging existing VNC password settings  ---------"
    rm -f $PASSWD_PATH
fi
VNC_PW_HASH=$(python3 -c "import crypt; print(crypt.crypt('${VNC_PW}', '\$5\$kasm\$'));")
#VNC_VIEW_PW_HASH=$(python3 -c "import crypt; print(crypt.crypt('${VNC_VIEW_ONLY_PW}', '\$5\$kasm\$'));")
echo "kasm_user:${VNC_PW_HASH}:ow" > $PASSWD_PATH
#echo "kasm_viewer:${VNC_VIEW_PW_HASH}:" >> $PASSWD_PATH
chmod 600 $PASSWD_PATH


# start processes
echo "Now starting KASM VNC..."
start_kasmvnc
#echo "Now starting window manager..."
#start_window_manager

# The following work only on KASM platform for now
#start_audio_out_websocket
#start_audio_out
#start_audio_in
#start_upload

STARTUP_COMPLETE=1


## log connect options
echo -e "\n\n------------------ KasmVNC environment started ------------------"

# tail vncserver logs
tail -f $HOME/.vnc/*$DISPLAY.log &

KASMIP=$(hostname -i)
echo "Kasm User ${KASM_USER}(${KASM_USER_ID}) started container id ${HOSTNAME} with local IP address ${KASMIP}"

# start custom startup script
custom_startup_script=/dockerstartup/custom_startup.sh
if [ -f "$custom_startup_script" ]; then
    if [ ! -x "$custom_startup_script" ]; then
        echo "${custom_startup_script}: not executable, exiting"
        exit 1
    fi

    "$custom_startup_script" &
fi

# Monitor Kasm Services
sleep 3
while :
do
    for process in "${!KASM_PROCS[@]}"; do
        if ! kill -0 "${KASM_PROCS[$process]}" ; then

            # If DLP Policy is set to fail secure, default is to be resilient
            if [[ ${DLP_PROCESS_FAIL_SECURE:-0} == 1 ]]; then
                exit 1
            fi

            case $process in
                kasmvnc)
                    echo "KasmVNC crashed, exiting container"
                    exit 1
                    # TODO: Is there a way to restore gracefully, restarting the container may be the best here
                    #start_kasmvnc
                    #/dockerstartup/custom_startup.sh
                    ;;
                kasm_audio_out_websocket)
                    echo "Restarting Audio Out Websocket Service"
                    start_audio_out_websocket
                    ;;
                kasm_audio_out)
                    echo "Restarting Audio Out Service"
                    start_audio_out
                    ;;
                kasm_audio_in)
                    echo "Audio In Service Failed"
                    # TODO: Needs work in python project to support auto restart
                    # start_audio_in
                    ;;
                upload_server)
                    echo "Restarting Upload Service"
                    # TODO: This will only work if both processes are killed, requires more work
                    start_upload
                    ;;
                *)
                    echo "Unknown Service: $process"
                    ;;
            esac
        fi
    done
    sleep 3
done


echo "Exiting Kasm container"