#!/usr/bin/env bash
set -e

if [[ $EUID -ne 0 ]]; then
   echo "This script must be run as root"
   exit 1
fi

TIMESTAMP=$(date +%s)
LOGFILE="./kasm_upgrade_${TIMESTAMP}.log"
function notify_err() {
  echo "An error has occurred please review the log at ${LOGFILE}"
}
function cleanup_log() {
    rm -f ${LOGFILE}
}
trap notify_err ERR
exec &> >(tee ${LOGFILE})

KASM_VERSION="1.19.0"
KASM_MAJOR_VERSION=$(echo ${KASM_VERSION} | awk -F'\\.' '{print $1}')
KASM_MINOR_VERSION=$(echo ${KASM_VERSION} | awk -F'\\.' '{print $2}')
CURRENT_VERSION=$(readlink -f /opt/kasm/current | awk -F'/' '{print $4}')
CURRENT_MAJOR_VERSION=$(echo ${CURRENT_VERSION} | awk -F'\\.' '{print $1}')
CURRENT_MINOR_VERSION=$(echo ${CURRENT_VERSION} | awk -F'\\.' '{print $2}')
KASM_INSTALL_BASE="/opt/kasm/${KASM_VERSION}"
OFFLINE_INSTALL="false"
DATABASE_HOSTNAME="false"
PUBLIC_HOSTNAME='false'
DB_PASSWORD="false"
DATABASE_PORT=5432
DATABASE_USER='kasmapp'
DATABASE_NAME='kasm'
REGISTRATION_TOKEN='false'
DEFAULT_GUAC_TOKEN='false'
GUAC_API_SERVER_HOSTNAME='false'
GUAC_PUBLIC_JWT_CERT='false'
DEFAULT_GUAC_ID='00000000-0000-0000-0000-000000000000'
SCRIPT_PATH="$( cd "$(dirname "$0")" ; pwd -P )"
KASM_RELEASE="$(realpath $SCRIPT_PATH)"
ARCH=$(uname -m)
DISK_SPACE=50000000000
DEFAULT_PROXY_LISTENING_PORT='443'
DEFAULT_RDP_LISTENING_PORT='3389'
API_SERVER_HOSTNAME='false'
ENABLE_LOSSLESS='false'
USE_ROLLING='false'
CHECK_DISK='true'
SKIP_EGRESS="false"
IGNORE_DEP_FAILURES="false"
ENABLE_EPEL=""
ARGS=("$@")

if [ "${ARCH}" != "x86_64" ] && [ "${ARCH}" != "aarch64" ] ; then
  if [ "${ARCH}" == "unknown" ] ; then
    echo "FATAL: Unable to detect system architecture"
    cleanup_log
    exit 1
  else
    echo "FATAL: Unsupported architecture"
    cleanup_log
    exit 1
  fi
fi

# Functions
function check_version_update_rules() {
  # check rules that are for specfic versions

  # 1.16.0
  if [ "${CURRENT_MINOR_VERSION}" -lt 16 ] && [ "${KASM_MINOR_VERSION}" -gt 15 ] && [ "${KASM_MAJOR_VERSION}" == 1 ] ; then
    if ([ "${ROLE}" == "guac" ]) && [ "${REGISTRATION_TOKEN}" == "false" ]; then
      echo "FATAL: option -k|--registration-token is required for upgrading from version 1.15.0 or lower to version 1.16.0 or higher. You can find the Component Registration Token in the Admin UI under Global Settings, Authentication section."
      exit 1
    fi
    #Check for 3389 in use. This is specific for 1.16.0 upgrade because we don't use the port previously
    if [[ "${ROLE}" == "guac" || -z "${ROLE}" ]]; then
      re='^[0-9]+$'
      if ! [[ ${DEFAULT_RDP_LISTENING_PORT} =~ $re ]] ; then
        echo "error: DEFAULT_RDP_LISTENING_PORT, (${DEFAULT_RDP_LISTENING_PORT}) is not an integer" >&2
        cleanup_log
        exit 1
      fi

      if ((${DEFAULT_RDP_LISTENING_PORT} <= 0 || ${DEFAULT_RDP_LISTENING_PORT} > 65535 )); then
        echo "error: DEFAULT_RDP_LISTENING_PORT, (${DEFAULT_RDP_LISTENING_PORT}) is not in the valid port range"
        cleanup_log
        exit 1
      fi

      echo "Checking if DEFAULT_RDP_LISTENING_PORT (${DEFAULT_RDP_LISTENING_PORT}) is free"
      if lsof -Pi :${DEFAULT_RDP_LISTENING_PORT} -sTCP:LISTEN  ; then
        echo "Port (${DEFAULT_RDP_LISTENING_PORT}) is in use. Installation cannot continue."
        cleanup_log
        exit -1
      else
        echo "Port (${DEFAULT_RDP_LISTENING_PORT}) is not in use."
      fi
    fi
  fi
}

function configure_systemd_services() {
  # configure systemd services when upgrading to systemd-enabled versions
  if can_use_systemd; then
    echo "configuring systemd services for kasm workspaces"

    mkdir -p /etc/systemd/system

    cp ${SCRIPT_PATH}/services/kasm.service /etc/systemd/system/kasm.service

    if { [ "${ROLE}" == "agent" ] || [ -z "${ROLE}" ]; } && [ "${SKIP_EGRESS}" != "true" ]; then
      cp ${SCRIPT_PATH}/services/kasm-network-plugin.service /etc/systemd/system/kasm-network-plugin.service
      systemctl enable kasm-network-plugin.service
    else
      # remove dependency on kasm-network-plugin if egress is skipped
      if [ -f /etc/systemd/system/kasm.service ]; then
        sed -i 's/ kasm-network-plugin.service//g' /etc/systemd/system/kasm.service
      fi
    fi

    systemctl enable kasm.service
  fi
}

function post_install_rules() {
  # apply rules post install that are for specific versions/scenarios

  # ensure healthcheck directories exist
  mkdir -p /opt/kasm/${KASM_VERSION}/healthcheck/kasm_api
  mkdir -p /opt/kasm/${KASM_VERSION}/healthcheck/kasm_agent
  mkdir -p /opt/kasm/${KASM_VERSION}/healthcheck/kasm_manager
  mkdir -p /opt/kasm/${KASM_VERSION}/healthcheck/rdp_gateway
  mkdir -p /opt/kasm/${KASM_VERSION}/healthcheck/rdp_https_gateway

  # 1.16.0
  # components rdp_gw and rdp_https_gw are new and need the registration token, for single server installs we can try to get this after install
  if [ "${CURRENT_MINOR_VERSION}" -lt 16 ] && [ "${KASM_MINOR_VERSION}" -gt 15 ] && [ "${KASM_MAJOR_VERSION}" == 1 ] ; then
    if [ -z "${ROLE}" ]; then
      wait_until_api_ready
      REGISTRATION_TOKEN=$(docker exec kasm_api python3 /src/api_server/server.pyc --cfg /opt/kasm/current/conf/app/api/api.app.config.yaml --get-global-setting auth:registration_token 2>/dev/null | grep -Po '(?<=auth:registration_token value: )\S+')
      if [ "${REGISTRATION_TOKEN}" != "false" ] && [ ! -z "${REGISTRATION_TOKEN}" ]; then
        echo "Update Rule (<=1.15.0 to >=1.16.0): Obtaining deployment's registration token and applying to new components."
        ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.rdp-gateway.registration_token = "'${REGISTRATION_TOKEN}'"' /opt/kasm/${KASM_VERSION}/conf/app/rdp_https_gateway/rdp_https_gateway.app.config.yaml
        ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.rdp-gateway.registration_token = "'${REGISTRATION_TOKEN}'"' /opt/kasm/${KASM_VERSION}/conf/app/rdp_gateway/passthrough.app.config.yaml
        chown kasm:kasm /opt/kasm/${KASM_VERSION}/conf/app/rdp_https_gateway/rdp_https_gateway.app.config.yaml
        chown kasm:kasm /opt/kasm/${KASM_VERSION}/conf/app/rdp_gateway/passthrough.app.config.yaml
        docker restart kasm_rdp_gateway kasm_rdp_https_gateway
      fi
    fi
  fi
  # 1.18.0
  # Share server and redis server removed in 1.18.0 so if the version upgraded from was less than 1.18.0 remove
  if [ "${CURRENT_MINOR_VERSION}" -lt 18 ] && [ "${KASM_MAJOR_VERSION}" == 1 ] ; then
    if docker ps -a --format '{{.Names}}' | grep -q "^kasm_share$"; then
      sudo docker rm -f kasm_share
    fi
    if docker ps -a --format '{{.Names}}' | grep -q "^kasm_redis$"; then
      sudo docker rm -f kasm_redis
    fi
  fi
}

function wait_until_api_ready() {
  local -r -i max_attempts=30
  local -i attempt_num=1
  echo "Waiting for API Server to become healthy"
  until docker inspect --format='{{json .State.Health}}' kasm_api | grep '"Status":"healthy"' > /dev/null 2>&1
  do
    if (( attempt_num == max_attempts ))
    then
      echo "API Server failed to enter healthy state"
      return 1
    else
      sleep $(( attempt_num++ ))
    fi
  done
  echo "API Server is healthy"
}

function gather_vars() {
  local agent_conf_name="agent.app.config.yaml"
  local api_conf_name="api.app.config.yaml"
  local guac_conf_name="kasmguac.app.config.yaml"
  local rdp_https_gw_conf_name="rdp_https_gateway.app.config.yaml"
  local rdp_gw_conf_name="passthrough.app.config.yaml"

  local app_conf_dir="/opt/kasm/${CURRENT_VERSION}/conf/app"
  # 1.17.0 Moved the app conf files to their own directories
  if [ -d "${app_conf_dir}/api" ] ; then
    local api_conf="${app_conf_dir}/api/${api_conf_name}"
    local guac_conf="${app_conf_dir}/guac/${guac_conf_name}"
    CURRENT_AGENT_CONF="${app_conf_dir}/agent/${agent_conf_name}"
    CURRENT_RDP_GATEWAY_PASSTHROUGH_CONF="${app_conf_dir}/rdp_gateway/${rdp_gw_conf_name}"
    CURRENT_RDP_HTTPS_GATEWAY_CONF="${app_conf_dir}/rdp_https_gateway/${rdp_https_gw_conf_name}"
  else
    local api_conf="${app_conf_dir}/${api_conf_name}"
    local guac_conf="${app_conf_dir}/${guac_conf_name}"
    CURRENT_AGENT_CONF="${app_conf_dir}/${agent_conf_name}"
    CURRENT_RDP_GATEWAY_PASSTHROUGH_CONF="${app_conf_dir}/${rdp_gw_conf_name}"
    CURRENT_RDP_HTTPS_GATEWAY_CONF="${app_conf_dir}/${rdp_https_gw_conf_name}"
  fi

  if [ "${DB_PASSWORD}" == "false" ] ; then
    DB_PASSWORD=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.database.password' "${api_conf}")
  fi
  MANAGER_TOKEN=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.manager.token' "${CURRENT_AGENT_CONF}")
  MANAGER_HOSTNAME=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.manager.hostnames.[0]' "${CURRENT_AGENT_CONF}")
  SERVER_ID=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.agent.server_id' "${CURRENT_AGENT_CONF}")
  PUBLIC_HOSTNAME=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.agent.public_hostname' "${CURRENT_AGENT_CONF}")
  MANAGER_ID=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.manager.manager_id' "${api_conf}")
  SERVER_HOSTNAME=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.server.server_hostname' "${api_conf}")
  if [ "${DATABASE_HOSTNAME}" == "false" ] ; then
    DATABASE_HOSTNAME=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.database.host' "${api_conf}")
  fi
  DB_IMAGE_NAME=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.services.db.image'  /opt/kasm/${CURRENT_VERSION}/docker/.conf/docker-compose-db.yaml)
  DATABASE_USER=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.database.username'  "${api_conf}")
  DATABASE_NAME=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.database.name'  "${api_conf}")
  if [ -f "${guac_conf}" ] && ([ "${ROLE}" == "guac" ] || [ -z "${ROLE}" ]); then
    DEFAULT_GUAC_TOKEN=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.api.token' "${guac_conf}")
    if [ -z "${DEFAULT_GUAC_TOKEN}" ] || [ "${DEFAULT_GUAC_TOKEN}" == "null" ]; then
      DEFAULT_GUAC_TOKEN=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.api.auth_token' "${guac_conf}")
    fi
    DEFAULT_GUAC_ID=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.kasmguac.id' "${guac_conf}")
    GUAC_API_SERVER_HOSTNAME=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.api.hostnames.[0]' "${guac_conf}")
    GUAC_PUBLIC_JWT_CERT=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.api.public_jwt_cert.[0]' "${guac_conf}")
    GUAC_PUBLIC_HOSTNAME=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.kasmguac.server_address' "${guac_conf}")
  fi
  if [ -f "${CURRENT_RDP_GATEWAY_PASSTHROUGH_CONF}" ] && ([ "${ROLE}" == "guac" ] || [ -z "${ROLE}" ]); then
    DEFAULT_RDP_GATEWAY_ID=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.rdp-gateway.id' "${CURRENT_RDP_GATEWAY_PASSTHROUGH_CONF}")
    RDP_GATEWAY_API_SERVER_HOSTNAME=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.api.hostnames.[0]' "${CURRENT_RDP_GATEWAY_PASSTHROUGH_CONF}")
    if [ -z "${DEFAULT_RDP_GATEWAY_TOKEN}" ]; then
      DEFAULT_RDP_GATEWAY_TOKEN=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.api.auth_token' "${CURRENT_RDP_GATEWAY_PASSTHROUGH_CONF}")
    fi
  fi
  if [ -f "${CURRENT_RDP_HTTPS_GATEWAY_CONF}" ] && ([ "${ROLE}" == "guac" ] || [ -z "${ROLE}" ]); then
    RDP_HTTPS_GW_ID=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.rdp-gateway.id' "${CURRENT_RDP_HTTPS_GATEWAY_CONF}")
    RDP_HTTPS_GW_PORT=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.rdp-gateway.server_port' "${CURRENT_RDP_HTTPS_GATEWAY_CONF}")
    RDP_HTTPS_GW_ADDRESS=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.rdp-gateway.server_address' "${CURRENT_RDP_HTTPS_GATEWAY_CONF}")
    RDP_HTTPS_GW_ZONE=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.rdp-gateway.zone' "${CURRENT_RDP_HTTPS_GATEWAY_CONF}")

    RDP_HTTPS_GW_API_TOKEN=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.api.auth_token' "${CURRENT_RDP_HTTPS_GATEWAY_CONF}")
    RDP_HTTPS_GW_API_ADDRESS=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.api.hostnames.[0]' "${CURRENT_RDP_HTTPS_GATEWAY_CONF}")
    RDP_HTTPS_GW_API_PORT=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.api.port' "${CURRENT_RDP_HTTPS_GATEWAY_CONF}")
  fi
  API_IMAGE_NAME=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.services.kasm_api.image'  /opt/kasm/${CURRENT_VERSION}/docker/.conf/docker-compose-all.yaml)
}

function can_use_systemd() {
  if [ -f "/opt/kasm/${KASM_VERSION}/services/kasm.service" ] || [ -f "${SCRIPT_PATH}/services/kasm.service" ]; then
    if command -v systemctl >/dev/null 2>&1; then
      return 0
    fi
  fi
  return 1
}

function is_kasm_service_enabled() {
  if systemctl is-enabled kasm.service >/dev/null 2>&1; then
    return 0
  fi
  return 1
}

function remove_duplicate_restarts() {
  ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i 'del(.services[].restart)' /opt/kasm/${CURRENT_VERSION}/docker/docker-compose.yaml
  ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.services[].restart = "unless-stopped"' /opt/kasm/${CURRENT_VERSION}/docker/docker-compose.yaml
  chown kasm:kasm /opt/kasm/${CURRENT_VERSION}/docker/docker-compose.yaml
  chmod 644 /opt/kasm/${CURRENT_VERSION}/docker/docker-compose.yaml
}

function stop_kasm() {
  if can_use_systemd && is_kasm_service_enabled; then
    systemctl stop kasm.service
  else
    /opt/kasm/bin/stop
  fi
}

function start_kasm() {
  if can_use_systemd && is_kasm_service_enabled; then
    systemctl start kasm.service
  else
    /opt/kasm/bin/start
  fi
}

function remove_known_orphans() {
  echo "Checking for known orphan containers (kasm_share, kasm_redis)"

  for container in kasm_share kasm_redis; do
    if docker ps -a --format '{{.Names}}' | grep -q "^${container}$"; then
      echo "Stopping and removing container: $container"
      docker stop "$container" || true
      docker rm -f "$container" || true
    else
      echo "Container $container does not exist. Skipping."
    fi
  done
}

function backup_db() {
  mkdir -p /opt/kasm/backups/
  if [ "${CURRENT_MAJOR_VERSION}" == "1" ] && [ "${CURRENT_MINOR_VERSION}" -ge "17" ] ; then
    # In 1.17.0 we changed the db user to postgres uid 70, so the backup folder and any backup files inside must also match that.
    chown -R 70:70 /opt/kasm/backups/
  fi
  # Hotfix in case we have old paths for dumping DB
  sed -i 's/tmp/backup/g' /opt/kasm/${CURRENT_VERSION}/bin/utils/db_backup
  # Run backup
  if [ "${ROLE}" == "db" ] || [ -z "${ROLE}" ] ; then
    bash /opt/kasm/${CURRENT_VERSION}/bin/utils/db_backup -f /opt/kasm/backups/kasm_db_backup.tar -p /opt/kasm/${CURRENT_VERSION}/
  else
    if [ "${CURRENT_MAJOR_VERSION}" == "1" ] && [ "${CURRENT_MINOR_VERSION}" -le "10" ] ; then
      # Kasm versions earlier than 1.11.0 didn't have a backup utility that worked with remote databases
      bash ${SCRIPT_PATH}/bin/utils/db_backup -f /opt/kasm/backups/kasm_db_backup.tar -p /opt/kasm/${CURRENT_VERSION}/ -q "${DATABASE_HOSTNAME}" -c ${DATABASE_USER} -j ${DATABASE_NAME}
    else
      bash /opt/kasm/${CURRENT_VERSION}/bin/utils/db_backup -f /opt/kasm/backups/kasm_db_backup.tar -p /opt/kasm/${CURRENT_VERSION}/ -q "${DATABASE_HOSTNAME}"
    fi
  fi
  if [ ! -f "/opt/kasm/backups/kasm_db_backup.tar" ]; then
    echo "Error backing up database, please follow the instructions at https://kasmweb.com/docs/latest/upgrade/single_server_upgrade.html to manually upgrade your installation"
    start_kasm
    exit 1
  fi

  chmod 750 /opt/kasm/backups 2>/dev/null || true
  chown 70:70 /opt/kasm/backups/kasm_db_backup.tar 2>/dev/null || true
  chmod 640 /opt/kasm/backups/kasm_db_backup.tar 2>/dev/null || true
}

function clean_install() {
  echo "Installing Kasm Workspaces ${KASM_VERSION}"
  bash ${SCRIPT_PATH}/install.sh -b -N -e -H -D ${OPTS} ${ROLLING} ${LOSSLESS} -Q "${DB_PASSWORD}" -M "${MANAGER_TOKEN}" -L ${DEFAULT_PROXY_LISTENING_PORT} -c "${DATABASE_USER}" -j "${DATABASE_NAME}"
}

function db_install() {
  echo "Installing Kasm Workspaces ${KASM_VERSION}"
  if [ "${ROLE}" == "db" ] ; then
    bash ${SCRIPT_PATH}/install.sh -e -H -D ${OPTS} ${ROLLING} -S db -Q "${DB_PASSWORD}" -L ${DEFAULT_PROXY_LISTENING_PORT} -c "${DATABASE_USER}" -j "${DATABASE_NAME}"
  else
    bash ${SCRIPT_PATH}/install.sh -N -e -H -D ${OPTS} ${ROLLING} -S init_remote_db -q "${DATABASE_HOSTNAME}" -Q "${DB_PASSWORD}" -L ${DEFAULT_PROXY_LISTENING_PORT} -g "${DATABASE_MASTER_USER}" -G "${DATABASE_MASTER_PASSWORD}" -c "${DATABASE_USER}" -j "${DATABASE_NAME}"
  fi
}

function agent_install() {
  echo "Installing Kasm Workspaces ${KASM_VERSION}"
  bash ${SCRIPT_PATH}/install.sh -b -N -e -H ${OPTS} ${ROLLING} -S agent -D -p "${PUBLIC_HOSTNAME}" -m "${MANAGER_HOSTNAME}" -M "${MANAGER_TOKEN}" -L ${DEFAULT_PROXY_LISTENING_PORT}
}

function app_install() {
  echo "Installing Kasm Workspaces ${KASM_VERSION}"
  bash ${SCRIPT_PATH}/install.sh -N -e -H ${OPTS} ${ROLLING} -S app -D -q "${DATABASE_HOSTNAME}" -Q "${DB_PASSWORD}" -L ${DEFAULT_PROXY_LISTENING_PORT} -c "${DATABASE_USER}" -j "${DATABASE_NAME}"
}

function proxy_install() {
  echo "Installing Kasm Workspaces ${KASM_VERSION}"
  if [ "${API_SERVER_HOSTNAME}" == "false" ]; then
    echo "FATAL: option -n|--api-hostname is required for the proxy role"
    cleanup_log
    exit 1
  fi
  bash ${SCRIPT_PATH}/install.sh -e -H ${OPTS} ${ROLLING} -S proxy -L ${DEFAULT_PROXY_LISTENING_PORT} -n "${API_SERVER_HOSTNAME}"
}

function guac_install() {
  echo "Installing Kasm Workspaces ${KASM_VERSION}"
  if [ ! -z "${SERVICE_IMAGE_TARFILE}" ]; then
    OPTS="-s ${SERVICE_IMAGE_TARFILE}"
  fi
  if [ "${GUAC_API_SERVER_HOSTNAME}" != "false" ]; then
    echo "Using API server hostname ${GUAC_API_SERVER_HOSTNAME} for connection proxy server upgrade"
  elif [ "${API_SERVER_HOSTNAME}" != "false" ] && [ "${GUAC_API_SERVER_HOSTNAME}" == "false" ]; then
    GUAC_API_SERVER_HOSTNAME=${API_SERVER_HOSTNAME}
    echo "Using API server hostname ${GUAC_API_SERVER_HOSTNAME} for connection proxy server upgrade"
  else
    echo "FATAL: option -n|--api-hostname is required for the guac role"
    exit 1
  fi
  if [ "${GUAC_PUBLIC_HOSTNAME}" != "false" ]; then
    echo "Using ${GUAC_PUBLIC_HOSTNAME} as the connection proxy server hostname"
  elif [ "${PUBLIC_HOSTNAME}" != "false" ] && [ "${GUAC_PUBLIC_HOSTNAME}" == "false" ]; then
    GUAC_PUBLIC_HOSTNAME=${PUBLIC_HOSTNAME}
    echo "Using ${GUAC_PUBLIC_HOSTNAME} as the connection proxy server hostname"
  else
    echo "FATAL: option -p|--public-hostname is required for the guac role"
    exit 1
  fi
  if [ "${REGISTRATION_TOKEN}" == "false" ]; then
    echo "FATAL: option -k|--registration-token is required for the guac role"
    exit 1
  fi
  bash ${SCRIPT_PATH}/install.sh -e -H ${OPTS} ${ROLLING} -S guac -L ${DEFAULT_PROXY_LISTENING_PORT} -n "${GUAC_API_SERVER_HOSTNAME}" -k "${REGISTRATION_TOKEN}" -p "${GUAC_PUBLIC_HOSTNAME}"
}

function modify_agent_configs() {
  local conf_file="/opt/kasm/${KASM_VERSION}/conf/app/agent/agent.app.config.yaml"
  ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.agent.server_id = "'${SERVER_ID}'"' "${conf_file}"
  ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.agent.public_hostname = "'${PUBLIC_HOSTNAME}'"' "${conf_file}"
  # There may be multiple manager hostnames we want to populate all of them.
  MANAGER_HOSTNAMES=$(${SCRIPT_PATH}/bin/utils/yq_${ARCH} -o j '.manager.hostnames' "${CURRENT_AGENT_CONF}")
  ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i ".manager.hostnames |= ${MANAGER_HOSTNAMES}" "${conf_file}"
  chown kasm:kasm "${conf_file}"
  chmod 644 "${conf_file}"
}

function modify_api_configs() {
  local conf_file="/opt/kasm/${KASM_VERSION}/conf/app/api/api.app.config.yaml"
  ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.manager.manager_id = "'${MANAGER_ID}'"' "${conf_file}"
  ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.server.server_hostname = "'${SERVER_HOSTNAME}'"' "${conf_file}"
  chown kasm:kasm "${conf_file}"
  chmod 644 "${conf_file}"
}

function modify_guac_configs() {
  local conf_file="/opt/kasm/${KASM_VERSION}/conf/app/guac/kasmguac.app.config.yaml"
  if [ "${CURRENT_MAJOR_VERSION}" == "1" ] && [ "${CURRENT_MINOR_VERSION}" -ge "12" ] ; then
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.kasmguac.id = "'${DEFAULT_GUAC_ID}'"' "${conf_file}"
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.api.auth_token = "'${DEFAULT_GUAC_TOKEN}'"' "${conf_file}"
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.api.public_jwt_cert = "'${GUAC_PUBLIC_JWT_CERT}'"' "${conf_file}"
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i 'del(.kasmguac.registration_token)' "${conf_file}"
    chown kasm:kasm "${conf_file}"
    chmod 644 "${conf_file}"
  fi
}

function modify_rdp_gateway_configs() {
  local conf_file="/opt/kasm/${KASM_VERSION}/conf/app/rdp_gateway/passthrough.app.config.yaml"
  if [ "${CURRENT_MAJOR_VERSION}" == "1" ] && [ "${CURRENT_MINOR_VERSION}" -ge "16" ] ; then
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.rdp-gateway.id = load("'"${CURRENT_RDP_GATEWAY_PASSTHROUGH_CONF}"'").rdp-gateway.id' "${conf_file}"
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.api.auth_token = load("'"${CURRENT_RDP_GATEWAY_PASSTHROUGH_CONF}"'").api.auth_token' "${conf_file}"
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.api.public_jwt_cert = load("'"${CURRENT_RDP_GATEWAY_PASSTHROUGH_CONF}"'").api.public_jwt_cert.[0]' "${conf_file}"
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i 'del(.rdp-gateway.registration_token)' "${conf_file}"
  fi

  # this component is new for 1.16 and will need to be registered
  if [ "${CURRENT_MAJOR_VERSION}" == "1" ] && [ "${CURRENT_MINOR_VERSION}" -le 15 ] && [ "${KASM_MINOR_VERSION}" -ge 16 ] && [ "${KASM_MAJOR_VERSION}" == 1 ]; then
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.rdp-gateway.registration_token = "'${REGISTRATION_TOKEN}'"' "${conf_file}"
  fi

  chown kasm:kasm "${conf_file}"
  chmod 644 "${conf_file}"
}

function modify_rdp_https_gateway_configs() {
  local conf_file="/opt/kasm/${KASM_VERSION}/conf/app/rdp_https_gateway/rdp_https_gateway.app.config.yaml"
  if [ "${CURRENT_MAJOR_VERSION}" == "1" ] && [ "${CURRENT_MINOR_VERSION}" -ge "16" ] ; then
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.rdp-gateway.id = "'${RDP_HTTPS_GW_ID}'"' "${conf_file}"
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.rdp-gateway.server_address = "'${RDP_HTTPS_GW_ADDRESS}'"' "${conf_file}"
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i ".rdp-gateway.server_port = ${RDP_HTTPS_GW_PORT}" "${conf_file}"
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.rdp-gateway.zone = "'${RDP_HTTPS_GW_ZONE}'"' "${conf_file}"
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i 'del(.rdp-gateway.registration_token)' "${conf_file}"

    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.api.auth_token = "'${RDP_HTTPS_GW_API_TOKEN}'"' "${conf_file}"
    # Get the public JWT Cert from the current installation
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.api.public_jwt_cert = load("'"${CURRENT_RDP_HTTPS_GATEWAY_CONF}"'").api.public_jwt_cert' "${conf_file}"
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.api.hostnames = [ "'${RDP_HTTPS_GW_API_ADDRESS}'" ]' "${conf_file}"
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i ".api.port = ${RDP_HTTPS_GW_API_PORT}" "${conf_file}"
  fi

  # this component is new for 1.16 and will need to be registered
  if [ "${CURRENT_MAJOR_VERSION}" == "1" ] && [ "${CURRENT_MINOR_VERSION}" -le 15 ] && [ "${KASM_MINOR_VERSION}" -ge 16 ] && [ "${KASM_MAJOR_VERSION}" == 1 ]; then
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} -i '.rdp-gateway.registration_token = "'${REGISTRATION_TOKEN}'"' "${conf_file}"
  fi

  chown kasm:kasm "${conf_file}"
  chmod 644 "${conf_file}"
}

function copy_nginx() {
  if [ -d "/opt/kasm/${CURRENT_VERSION}/conf/nginx/containers.d" ] && [ "$(ls -A /opt/kasm/${CURRENT_VERSION}/conf/nginx/containers.d 2>/dev/null)" ]; then
    cp /opt/kasm/${CURRENT_VERSION}/conf/nginx/containers.d/* /opt/kasm/${KASM_VERSION}/conf/nginx/containers.d/
    chown root:root /opt/kasm/${KASM_VERSION}/conf/nginx/containers.d/*
    chmod 644 /opt/kasm/${KASM_VERSION}/conf/nginx/containers.d/*
  fi
}

function restore_db() {

  if [ ! -r "/opt/kasm/backups/kasm_db_backup.tar" ]; then
    echo "FATAL: /opt/kasm/backups/kasm_db_backup.tar is not readable. Check permissions/ownership (umask, chmod/chown)."
    exit 1
  fi

  if [ "${ROLE}" == "db" ] || [ -z "${ROLE}" ] ; then
    /opt/kasm/${KASM_VERSION}/bin/utils/db_restore -a -f /opt/kasm/backups/kasm_db_backup.tar -p  /opt/kasm/${KASM_VERSION} -c ${DATABASE_USER} -j ${DATABASE_NAME}
    /opt/kasm/${KASM_VERSION}/bin/utils/db_upgrade -p /opt/kasm/${KASM_VERSION}
  else
    /opt/kasm/${KASM_VERSION}/bin/utils/db_restore -a -f /opt/kasm/backups/kasm_db_backup.tar -p  /opt/kasm/${KASM_VERSION} -q "${DATABASE_HOSTNAME}" -g "${DATABASE_MASTER_USER}" -G "${DATABASE_MASTER_PASSWORD}" -c ${DATABASE_USER} -j ${DATABASE_NAME}
    /opt/kasm/${KASM_VERSION}/bin/utils/db_upgrade -p /opt/kasm/${KASM_VERSION} -q "${DATABASE_HOSTNAME}"
  fi
}

function copy_certificates() {
  # We check for our self signed certificates, and if those are present and have less than 6 months left
  # (two Kasm releases at our current quarterly release cycle) we generate new certificates.
  subject=$(openssl x509 -in /opt/kasm/${CURRENT_VERSION}/certs/kasm_nginx.crt -noout -subject)
  if [[ ${subject} == *"OU=DoFu"* ]] && [[ ${subject} == *"emailAddress=none@none.none"* ]]; then
    openssl x509 -in /opt/kasm/${CURRENT_VERSION}/certs/kasm_nginx.crt -noout -checkend 15768017
    ret_code=$?
    if [ $ret_code -ne 0 ] ; then
      echo "Existing self-signed certs expire within six months generating new self-signed certs"
      sudo openssl req -x509 -nodes -days 1825 -newkey rsa:2048 -keyout ${KASM_INSTALL_BASE}/certs/kasm_nginx.key -out ${KASM_INSTALL_BASE}/certs/kasm_nginx.crt -subj "/C=US/ST=VA/L=None/O=None/OU=DoFu/CN=$(hostname)/emailAddress=none@none.none" 2> /dev/null
      return
    fi
  fi
  echo "Copying existing certs from Kasm Workspaces ${CURRENT_VERSION} install"
  cp -f /opt/kasm/${CURRENT_VERSION}/certs/kasm_nginx.crt ${KASM_INSTALL_BASE}/certs/kasm_nginx.crt
  cp -f /opt/kasm/${CURRENT_VERSION}/certs/kasm_nginx.key ${KASM_INSTALL_BASE}/certs/kasm_nginx.key
  chown kasm:kasm ${KASM_INSTALL_BASE}/certs/kasm_nginx.key
  chmod 600 ${KASM_INSTALL_BASE}/certs/kasm_nginx.key
  chown kasm:kasm ${KASM_INSTALL_BASE}/certs/kasm_nginx.crt
  chmod 600 ${KASM_INSTALL_BASE}/certs/kasm_nginx.crt
}

function display_help() {
  CMD='\033[0;31m'
  NC='\033[0m'
  echo "Kasm Upgrader ${KASM_VERSION}" 
  echo "Usage IE:"
  echo "${0} --add-images --proxy-port 443"
  echo    ""
  echo    "Flag                                        Description"
  echo    "---------------------------------------------------------------------------------------------------------------"
  echo -e "| ${CMD}-h|--help${NC}                     | Display this help menu                                                      |"
  echo -e "| ${CMD}-L|--proxy-port${NC}               | Default Proxy Listening Port                                                |"
  echo -e "| ${CMD}-s|--offline-service${NC}          | Path to the tar.gz service images offline installer                         |"
  echo -e "| ${CMD}-x|--offline-network-plugin${NC}   | Path to the tar.gz network plugin offline installer                         |"
  echo -e "| ${CMD}--offline-logger-plugin${NC}       | Path to the tar.gz logger plugin offline installer                          |"
  echo -e "| ${CMD}-S|--role${NC}                     | Role to Upgrade: [app|db|agent|remote_db|guac|proxy]                        |"
  echo -e "| ${CMD}-p|--public-hostname${NC}          | Agent/Component <IP/Hostname> used to register with deployment.             |"
  echo -e "| ${CMD}-g|--database-master-user${NC}     | Database master username required for remote DB                             |"
  echo -e "| ${CMD}-G|--database-master-password${NC} | Database master password required for remote DB                             |"
  echo -e "| ${CMD}-q|--db-hostname${NC}              | Database Hostname needed when upgrading agent and pulling images            |"
  echo -e "| ${CMD}-T|--db-port${NC}                  | Database port needed when upgrading agent and pulling images (default 5432) |"
  echo -e "| ${CMD}-Q|--db-password${NC}              | Database Password needed when upgrading agent and pulling images            |"
  echo -e "| ${CMD}-b|--no-check-disk${NC}            | Do not check disk space                                                     |"
  echo -e "| ${CMD}-n|--api-hostname${NC}             | Set API server hostname                                                     |"
  echo -e "| ${CMD}-A|--enable-lossless${NC}          | Enable lossless streaming option (1.12 and above)                           |"
  echo -e "| ${CMD}-O|--use-rolling-images${NC}       | Use rolling Service images                                                  |"
  echo -e "| ${CMD}-f|--use-static-images${NC}        | Use static versioned Service images                                         |"
  echo -e "| ${CMD}-k|--registration-token${NC}       | Register a component with an existing deployment.                           |"
  echo -e "| ${CMD}--skip-egress${NC}                 | Do not install egress network plugin or dependencies                        |"
  echo -e "| ${CMD}--ignore-dep-failures${NC}         | Continue installation even if optional dependencies fail to install         |"
  echo -e "| ${CMD}--enable-epel${NC}                 | Enable EPEL repo to install packages not in default RHEL repos              |"
  echo    "---------------------------------------------------------------------------------------------------------------"
}

function check_role() {
if [ "${ROLE}" != "agent" ] && [ "${ROLE}" !=  "app" ] && [ "${ROLE}" != "db" ] && [ "${ROLE}" != "remote_db" ] && [ "${ROLE}" != "guac" ] &&  [ "${ROLE}" != "proxy" ];
then
  echo "Invalid Role Defined"
  display_help
  cleanup_log
  exit 1
fi
}

# Command line opts
for index in "${!ARGS[@]}"; do
  case ${ARGS[index]} in
    -L|--proxy-port)
      DEFAULT_PROXY_LISTENING_PORT="${ARGS[index+1]}"
      echo "Setting Default Listening Port as ${DEFAULT_PROXY_LISTENING_PORT}"
      ;;
    -S|--role)
      ROLE="${ARGS[index+1]}"
      check_role
      echo "Setting Role as ${ROLE}"
      ;;
    -h|--help)
      display_help
      cleanup_log
      exit 0
      ;;
    -s|--offline-service)
      SERVICE_IMAGE_TARFILE="${ARGS[index+1]}"
      OFFLINE_INSTALL="true"

      if [ ! -f "$SERVICE_IMAGE_TARFILE" ]; then
        echo "FATAL: Service image tarfile does not exist: ${SERVICE_IMAGE_TARFILE}"
        cleanup_log
        exit 1
      fi

      echo "Setting service image tarfile to ${SERVICE_IMAGE_TARFILE}"
      ;;
    -x|--offline-network-plugin)
      NETWORK_PLUGIN_IMAGE_TARFILE="${ARGS[index+1]}"

        if [ ! -f "$NETWORK_PLUGIN_IMAGE_TARFILE" ]; then
          echo "FATAL: Network plugin image tarfile does not exist: ${NETWORK_PLUGIN_IMAGE_TARFILE}"
          cleanup_log
          exit 1
        fi

        echo "Setting network plugin image tarfile to ${NETWORK_PLUGIN_IMAGE_TARFILE}"
      ;;
    --offline-logger-plugin)
      LOGGER_PLUGIN_IMAGE_TARFILE="${ARGS[index+1]}"
      if [ ! -f "$LOGGER_PLUGIN_IMAGE_TARFILE" ]; then
        echo "FATAL: Logger plugin image tarfile does not exist: ${LOGGER_PLUGIN_IMAGE_TARFILE}"
        cleanup_log
        exit 1
      fi
      echo "Setting logger plugin image tarfile to ${LOGGER_PLUGIN_IMAGE_TARFILE}"
      ;;
    -g|--database-master-user)
      DATABASE_MASTER_USER="${ARGS[index+1]}"
      echo "Using Database Master User ${DATABASE_MASTER_USER}"
      ;;
    -G|--database-master-password)
      DATABASE_MASTER_PASSWORD="${ARGS[index+1]}"
      echo "Using Database Master Password from stdin -G"
      ;;
    -q|--db-hostname)
      DATABASE_HOSTNAME="${ARGS[index+1]}"
      echo "Setting Database Hostname as ${DATABASE_HOSTNAME}"
      ;;
    -T|--db-port)
      DATABASE_PORT="${ARGS[index+1]}"
      echo "Setting Database Port to ${DATABASE_PORT}"
      ;;
    -Q|--db-password)
      DB_PASSWORD="${ARGS[index+1]}"
      echo "Setting Default Database Password from stdin -Q"
      ;;
    -n|--api-hostname)
      API_SERVER_HOSTNAME="${ARGS[index+1]}"
      echo "Setting API Server Hostname as ${API_SERVER_HOSTNAME}"
      ;;
    -A|--enable-lossless)
      ENABLE_LOSSLESS="true"
      ;;
    -O|--use-rolling-images)
      USE_ROLLING="true"
      ;;
    -l|--use-staging-images)
      USE_STAGING="true"
      ;;
    -f|--use-static-images)
      USE_STATIC="true"
      ;;
    -k|--registration-token)
      REGISTRATION_TOKEN="${ARGS[index+1]}"
      ;;
    -b|--no-check-disk)
      CHECK_DISK="false"
      ;;
    -p|--public-hostname)
      PUBLIC_HOSTNAME="${ARGS[index+1]}"
      echo "Setting Public Hostname as ${PUBLIC_HOSTNAME}"
      ;;
    --skip-egress)
      SKIP_EGRESS="true"
      ;;
    --ignore-dep-failures)
      IGNORE_DEP_FAILURES="true"
      ;;
    --enable-epel)
      ENABLE_EPEL="true"
      ;;
    -*|--*)
      echo "Unknown option ${ARGS[index]}"
      display_help
      cleanup_log
      exit 1
      ;;
  esac
done

# Set some installer variables based on flags passed
LOSSLESS=""
if [ "${ENABLE_LOSSLESS}" == "true" ]; then
  LOSSLESS="-A"
fi
ROLLING=""
if [ "${USE_ROLLING}" == "true" ]; then
  ROLLING="-O"
fi
if [ "${USE_STAGING}" == "true" ]; then
  ROLLING="-l"
fi
if [ "${USE_STATIC}" == "true" ]; then
  ROLLING="-f"
fi
if [ ! -z "${SERVICE_IMAGE_TARFILE}" ]; then
  OPTS="-s ${SERVICE_IMAGE_TARFILE} -I"
else
  OPTS="-I"
fi
if [ ! -z "${NETWORK_PLUGIN_IMAGE_TARFILE}" ]; then
  OPTS="${OPTS} -x ${NETWORK_PLUGIN_IMAGE_TARFILE}"
fi
if [ ! -z "${LOGGER_PLUGIN_IMAGE_TARFILE}" ]; then
  OPTS="${OPTS} --offline-logger-plugin ${LOGGER_PLUGIN_IMAGE_TARFILE}"
fi
if [ "${SKIP_EGRESS}" == "true" ]; then
  OPTS="${OPTS} --skip-egress"
fi
if [ "${IGNORE_DEP_FAILURES}" == "true" ]; then
  OPTS="${OPTS} --ignore-dep-failures"
fi
if [ "${ENABLE_EPEL}" == "true" ]; then
  OPTS="${OPTS} --enable-epel"
fi

check_version_update_rules

function get_db_alembic_version() {
  local db_ver
  db_ver=$(docker exec kasm_db psql -U kasmapp -d kasm -tAc "select version_num from alembic_version;" 2>/dev/null)

  if [ -n "$db_ver" ]; then
      echo "$db_ver"
      return 0
  fi

  # fallback in case DB is not running 
  if [ -f "/opt/kasm/current/conf/database/seed_data/default_agents.yaml" ]; then
      ${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.alembic_version' \
        /opt/kasm/current/conf/database/seed_data/default_agents.yaml
      return 0
  fi

  echo ""
  return 1
}

function get_installer_alembic_version() {
    ${SCRIPT_PATH}/bin/utils/yq_${ARCH} '.alembic_version' \
      ${KASM_INSTALL_BASE}/conf/database/seed_data/default_agents.yaml
}

if [ "${CURRENT_VERSION}" == "${KASM_VERSION}" ]; then
  
  DB_ALEMBIC_VERSION=$(get_db_alembic_version)
  INSTALLER_ALEMBIC_VERSION=$(get_installer_alembic_version)
  
  if [[ "${DB_ALEMBIC_VERSION}" == "${INSTALLER_ALEMBIC_VERSION}"* ]]; then
    echo "No upgrade necessary. Kasm Workspaces ${CURRENT_VERSION} with schema ${DB_ALEMBIC_VERSION} is already up to date."
    cleanup_log
    exit 0
  fi
fi

# Perform upgrade
if [ -z "${ROLE}" ]; then
  gather_vars
  remove_duplicate_restarts
  stop_kasm
  remove_known_orphans
  backup_db
  clean_install
  modify_agent_configs
  modify_api_configs
  modify_guac_configs
  modify_rdp_gateway_configs
  modify_rdp_https_gateway_configs
  copy_certificates
  copy_nginx
  restore_db
  configure_systemd_services
  start_kasm
elif [ "${ROLE}" == "db" ]; then
  gather_vars
  remove_duplicate_restarts
  stop_kasm
  remove_known_orphans
  backup_db
  db_install
  restore_db
  configure_systemd_services
  start_kasm
elif [ "${ROLE}" == "remote_db" ]; then
  gather_vars
  remove_duplicate_restarts
  stop_kasm
  remove_known_orphans
  backup_db
  db_install
  restore_db
elif [ "${ROLE}" == "agent" ]; then
  gather_vars
  remove_duplicate_restarts
  stop_kasm
  remove_known_orphans
  agent_install
  modify_agent_configs
  copy_certificates
  copy_nginx
  configure_systemd_services
  start_kasm
elif [ "${ROLE}" == "app" ]; then
  gather_vars
  remove_duplicate_restarts
  stop_kasm
  remove_known_orphans
  app_install
  modify_api_configs
  copy_certificates
  configure_systemd_services
  start_kasm
elif [ "${ROLE}" == "guac" ]; then
  gather_vars
  remove_duplicate_restarts
  stop_kasm
  remove_known_orphans
  guac_install
  modify_guac_configs
  modify_rdp_gateway_configs
  modify_rdp_https_gateway_configs
  copy_certificates
  configure_systemd_services
  start_kasm
elif [ "${ROLE}" == "proxy" ]; then
  remove_duplicate_restarts
  stop_kasm
  remove_known_orphans
  proxy_install
  copy_certificates
  configure_systemd_services
  start_kasm
fi

post_install_rules

printf "\n\n"
echo "Upgrade from ${CURRENT_VERSION} to ${KASM_VERSION} Complete"

cleanup_log
