#!/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 $1 $2
}

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
}

getAwsRoleArns() {
  /usr/lib/rstudio-server/bin/rserver-db get-aws-role-arns "$1"
}

setAwsRoleArns() {
  /usr/lib/rstudio-server/bin/rserver-db set-aws-role-arns "$1" --role-arns="$2"
}

stopNode() {
  /usr/lib/rstudio-server/bin/rserver-db stop-node
}

listNodes() {
  /usr/lib/rstudio-server/bin/rserver-db list-nodes
}

resetCluster() {
  /usr/lib/rstudio-server/bin/rserver-db reset-cluster
}

deleteNode() {
  /usr/lib/rstudio-server/bin/rserver-db delete-node $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
}

configureVSCode() {
  /usr/lib/rstudio-server/extras/scripts/configure-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)
        if isServerRunning ; then stopNode; fi
        daemonCmd "stop"
        ;;

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

    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"
        ;;

    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 "2023.03.0+386.pro1 Workbench (Cherry Blossom) for Ubuntu Jammy"
        ;;

    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 "$2" "$3"
        ;;

    list-users)
        listUsers
        ;;

    add-user)
        addUser $2 $3
        ;;

    set-admin)
        setAdmin $2 $3
        ;;

    lock-user)
        lockUser $2
        ;;

    unlock-user)
        unlockUser $2
        ;;

    get-aws-role-arns)
        getAwsRoleArns "$2"
        ;;
    
    set-aws-role-arns)
        setAwsRoleArns "$2" "$3"
        ;;
    
    list-nodes)
       listNodes
       ;;

    reset-cluster)
       resetCluster
       ;;

    delete-node)
       if [ -n "$2" ]; then
          while true; do
             read -p "
WARNING This command is only intended to be used to delete offline nodes without any active sessions, but does NOT verify this state before deleting the node. If node '$2' has active sessions, data loss may occur within those sessions.

To verify the node's status run 'rstudio-server list-nodes' or 'curl http://localhost:8787/load-balancer/status'.

Run 'rstudio-server stop' from the to-be-deleted node or 'man rstudio-server' to view options for suspending sessions.

Are you sure you want to delete node '$2' (y or n)? " yn
             case $yn in
                [Yy]* ) deleteNode $2; break;;
                [Nn]* ) break;;
                * ) echo "Invalid input - please enter yes or no";;
             esac
          done
       else
          echo "Must specify ID of node to delete."
       fi
       ;;
      
    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)
           echo "The install-vs-code command is deprecated, use configure-vs-code instead"
           exit 1
           ;;
 
    configure-vs-code)
        configureVSCode "${@:2}"
        ;;

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

    *)
        echo $"Usage: rstudio-server {status|start|stop|restart|test-config|verify-installation|run-diagnostics|reload|configure-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|list-nodes|delete-node|reset-cluster}"
        echo ""
        echo "See the Posit Workbench Administration Guide or run 'man rstudio-server' for documentation on supported commands."
        exit 2
esac
