#!/bin/bash

daemonCmd() {
  INIT_SYSTEM=`cat /proc/1/comm 2>/dev/null`
  if test "$INIT_SYSTEM" = "systemd"
  then
     eval "systemctl $1 rstudio-server.service"
  elif test -e /etc/init/rstudio-server.conf
  then
     eval "initctl $1 rstudio-server"
  else
     eval "/etc/init.d/rstudio-server $1"
  fi
  return $?
}

testConfig() {
  `/usr/lib/rstudio-server/bin/rserver --test-config`
}

runDiagnostics() {
  /usr/lib/rstudio-server/bin/run-diagnostics
}

verifyInstallation() {
  /usr/lib/rstudio-server/bin/rserver --verify-installation=1 --server-daemonize=0 $1 $2 $3 $4 $5 $6 $7 $8
}

listUsers() {
  /usr/lib/rstudio-server/bin/rserver-db list-users
}

addUser() {
  isAdmin=$2
  if [ -z "$isAdmin" ]; then
    isAdmin="0"
  fi
  /usr/lib/rstudio-server/bin/rserver-db add-user $1 --is-admin=$isAdmin
}

setAdmin() {
  /usr/lib/rstudio-server/bin/rserver-db set-admin $1 --is-admin=$2
}

lockUser() {
  /usr/lib/rstudio-server/bin/rserver-db lock-user $1
}

unlockUser() {
  /usr/lib/rstudio-server/bin/rserver-db unlock-user $1
}

encryptPass() {
  secureCookieKey=$(sed -n /^secure-cookie-key-file=/s/^secure-cookie-key-file=//p /etc/rstudio/rserver.conf)
  if [ -z "$secureCookieKey" ]; then
    secureCookieKey=""
  fi
  /usr/lib/rstudio-server/bin/rserver-encrypt "$secureCookieKey"
}

migrateDb() {
  /usr/lib/rstudio-server/bin/migrate.sh
}

installVSCode() {
  /usr/lib/rstudio-server/extras/scripts/install-vs-code "$@"
}

installVSCodeExt() {
  /usr/lib/rstudio-server/extras/scripts/install-vs-code-ext.sh "$@"
}

serverUser() {
  userStr=$(sed -n /^server-user=/s/^server-user=//p /etc/rstudio/rserver.conf)
  if [ -z "$userStr" ]; then
    userStr="rstudio-server"
  fi

  echo "$userStr"
}

launcherSessionsEnabled() {
  grep ^launcher-sessions-enabled=1 /etc/rstudio/rserver.conf &> /dev/null
}

killLauncherSession() {
  echo "Killing launcher session $1"
  output=$(/usr/lib/rstudio-server/bin/rpcsend "job_launcher_rpc/kill_session" "[\"$1\", \"$(serverUser)\"]")
  if [ $? -eq 0 ]; then
    (echo $output | grep '"result":true') &> /dev/null
    if [ $? -eq 0 ]; then
      echo "Kill completed successfully"
    else
      echo "Session $1 not found"
    fi
  else
    echo "$output"
  fi
}

suspendLauncherSession() {
  echo "Suspending launcher session $1"
  output=$(/usr/lib/rstudio-server/bin/rpcsend "job_launcher_session_rpc/$(serverUser)/$1/suspend_session" "[false]")
  if [ $? -eq 0 ]; then
    (echo $output | grep "error") &> /dev/null
    if [ $? -eq 1 ]; then
      echo "Suspend completed successfully"
    else
      echo "Session $1 not found"
    fi
  else
    echo "$output"
  fi
}

forceSuspendLauncherSession() {
  echo "Force suspending launcher session $1"
  output=$(/usr/lib/rstudio-server/bin/rpcsend "job_launcher_session_rpc/$(serverUser)/$1/suspend_session" "[true]")
  if [ $? -eq 0 ]; then
    (echo $output | grep "error") &> /dev/null
    if [ $? -eq 1 ]; then
      echo "Force suspend completed successfully"
    else
      echo "Session $1 not found"
    fi
  else
    echo "$output"
  fi
}

suspendAllLauncherSessions() {
  echo "Suspending all launcher sessions"
  output=$(/usr/lib/rstudio-server/bin/rpcsend "job_launcher_rpc/suspend_all_sessions" "[false, \"$(serverUser)\"]")
  if [ $? -eq 0 ]; then
    (echo $output | grep '"result":true') &> /dev/null
    if [ $? -eq 0 ]; then
      echo "Suspend all completed successfully"
    else
      echo "Suspend all did not complete successfully: $output"
    fi
  else
    echo "$output"
  fi

}

forceSuspendAllLauncherSessions() {
  echo "Force suspending all launcher sessions"
  output=$(/usr/lib/rstudio-server/bin/rpcsend "job_launcher_rpc/suspend_all_sessions" "[true, \"$(serverUser)\"]")
  if [ $? -eq 0 ]; then
    (echo $output | grep '"result":true') &> /dev/null
    if [ $? -eq 0 ]; then
      echo "Force suspend all completed successfully"
    else
      echo "Force suspend all did not complete successfully: $output"
    fi
  else
    echo "$output"
  fi

}

killAllLauncherSessions() {
  echo "Killing all launcher sessions"
  output=$(/usr/lib/rstudio-server/bin/rpcsend "job_launcher_rpc/kill_all_sessions" "[\"$(serverUser)\"]")
  if [ $? -eq 0 ]; then
    (echo $output | grep '"result":true') &> /dev/null
    if [ $? -eq 0 ]; then
      echo "Kill all completed successfully"
    else
      echo "Kill all did not complete successfully: $output"
    fi
  else
    echo "$output"
  fi
}

isServerRunning() {
 pgrep -x rserver &> /dev/null
}

case "$1" in

    status)
        daemonCmd "status"
        ;;

    start)
        testConfig || exit $?
        daemonCmd "start"
        ;;

    stop)
        daemonCmd "stop"
        if (pidof rsession >/dev/null); then killall -s USR2 $2 $3 $4 $5 $6 $7 $8 $9 rsession; fi
        if (pidof rworkspaces >/dev/null); then killall -s USR2 $2 $3 $4 $5 $6 $7 $8 $9 rworkspaces; fi
        ;;

    restart)
        testConfig || exit $?
        daemonCmd "restart"
        if (pidof rsession >/dev/null); then killall -s USR2 $2 $3 $4 $5 $6 $7 $8 $9 rsession; fi
        if (pidof rworkspaces >/dev/null); then killall -s USR2 $2 $3 $4 $5 $6 $7 $8 $9 rworkspaces; fi
        ;;

    test-config)
        testConfig
        ;;

    verify-installation)
        if isServerRunning ; then
           echo "Server is running and must be stopped before running verify-installation"
           exit 1
        fi

        # we call stop here even though the service should already be stopped
        # just in case (we don't want the service running when invoking verify-installation)
        daemonCmd "stop" 2>/dev/null

        verifyInstallation $2 $3 $4 $5 $6 $7 $8 $9
        ;;

    suspend-session)
        if [ -n "$2" ]; then
           if launcherSessionsEnabled; then
              suspendLauncherSession $2
           else
              kill -s USR1 $3 $4 $5 $6 $7 $8 $9 $2
           fi
        else
           if launcherSessionsEnabled; then
              argType="ID"
           else
              argType="PID"
           fi
           echo "Must specify $argType of session to suspend"
           exit 1
        fi
        ;;

    suspend-all)
        if (pidof rworkspaces >/dev/null); then killall -s USR1 $2 $3 $4 $5 $6 $7 $8 $9 rworkspaces; fi
        if launcherSessionsEnabled; then
           suspendAllLauncherSessions
        else
           if (pidof rsession >/dev/null); then killall -s USR1 $2 $3 $4 $5 $6 $7 $8 $9 rsession; fi
        fi
        ;;

    force-suspend-session)
        if [ -n "$2" ]; then
           if launcherSessionsEnabled; then
              forceSuspendLauncherSession $2
           else
              kill -s USR2 $3 $4 $5 $6 $7 $8 $9 $2
           fi
        else
           if launcherSessionsEnabled; then
              argType="ID"
           else
              argType="PID"
           fi
           echo "Must specify $argType of session to force suspend"
           exit 1
        fi
        ;;

    force-suspend-all)
        if (pidof rworkspaces >/dev/null); then killall -s USR2 $2 $3 $4 $5 $6 $7 $8 $9 rworkspaces; fi
        if launcherSessionsEnabled; then
           forceSuspendAllLauncherSessions
        else
           if (pidof rsession >/dev/null); then killall -s USR2 $2 $3 $4 $5 $6 $7 $8 $9 rsession; fi
        fi
        ;;

    kill-session)
        if [ -n "$2" ]; then
           if launcherSessionsEnabled; then
              killLauncherSession $2
           else
              kill -s KILL $3 $4 $5 $6 $7 $8 $9 $2
           fi
        else
           if launcherSessionsEnabled; then
              argType="ID"
           else
              argType="PID"
           fi
           echo "Must specify $argType of session to kill"
           exit 1
        fi
        ;;

    kill-all)
        if (pidof rworkspaces >/dev/null); then killall -s KILL $2 $3 $4 $5 $6 $7 $8 $9 rworkspaces; fi
        if launcherSessionsEnabled; then
           killAllLauncherSessions
        else
           if (pidof rsession >/dev/null); then killall -s KILL $2 $3 $4 $5 $6 $7 $8 $9 rsession; fi
        fi
        ;;

    upload-minidump)
        /usr/lib/rstudio-server/bin/crashpad_http_upload -f upload_file_minidump=$2 -u https://sentry.io/api/1379214/minidump/?sentry_key=4e76b8d2cffb49419fec1b431e09247c
        ;;

    offline)
        mkdir -p /var/lib/rstudio-server
        touch /var/lib/rstudio-server/offline
        daemonCmd "restart"
        if (pidof rsession >/dev/null); then killall -s USR2 $2 $3 $4 $5 $6 $7 $8 $9 rsession; fi
        if (pidof rworkspaces >/dev/null); then killall -s USR2 $2 $3 $4 $5 $6 $7 $8 $9 rworkspaces; fi
        ;;

    online)
        mkdir -p /var/lib/rstudio-server
        rm -f /var/lib/rstudio-server/offline
        daemonCmd "restart"
        ;;

    active-sessions)
        ps opid,cputime,args -C "rsession"
        ;;

    version)
        echo "1.4.1743-4 Workbench (Evergreen Rose) for Ubuntu Bionic"
        ;;

    license-manager)
        shift
        /usr/lib/rstudio-server/bin/license-manager "$@"

        # Save return value of licensing operation
        licenseResult=$?

        # If the server is currently running and the operation was successful,
        # then send it a SIGHUP to pick up the new licensing state (don't bother
        # in the case of a status/verify operation which doesn't change state)
        if [ "$licenseResult" == "0" ] && [ isServerRunning ] ; then 
            if [ "$1" != "status" ] && [ "$1" != "verify" ]; then
                kill -s SIGHUP `pidof rserver`
            fi
        fi

        # Return result of licensing operation
        exit $licenseResult
        ;;

    reload)
        if isServerRunning ; then
            PID=`pidof rserver`
            echo "Reloading supported server configuration values (process $PID)"
            kill -s SIGHUP $PID
        else
            echo "Configuration not reloaded (server not running)"
        fi
        ;;

    run-diagnostics)
        runDiagnostics
        ;;

    list-users)
        listUsers
        ;;

    add-user)
        addUser $2 $3
        ;;

    set-admin)
        setAdmin $2 $3
        ;;

    lock-user)
        lockUser $2
        ;;

    unlock-user)
        unlockUser $2
        ;;

    encrypt-password)
        if [[ $EUID -ne 0 ]]; then
          echo "encrypt-password must be executed as root"
          exit 1
        fi
        encryptPass
        ;;

    migrate-db)
        migrateDb
        ;;

    install-vs-code)
        installVSCode "${@:2}"
        ;;

    install-vs-code-ext)
        installVSCodeExt "${@:2}"
        ;;

    *)
        echo $"Usage: rstudio-server {status|start|stop|restart|test-config|verify-installation|run-diagnostics|reload|install-vs-code|install-vs-code-ext|suspend-session|suspend-all|force-suspend-session|force-suspend-all|kill-session|kill-all|upload-minidump|offline|online|active-sessions|version|list-users|add-user|set-admin|lock-user|unlock-user}"
        exit 2
esac
