source/idea/idea-bootstrap/virtual-desktop-host-linux/configure.sh.jinja2

#!/bin/bash
#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
#  with the License. A copy of the License is located at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
#  OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
#  and limitations under the License.

set -x

SEMAPHORE_DIR="/root/bootstrap/semaphore"
CONFIG_FINISHED_LOCK="${SEMAPHORE_DIR}/configure_finished.lock"
INSTANCE_READY_LOCK="${SEMAPHORE_DIR}/instance_ready.lock"
SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )

if [[ ! -f ${CONFIG_FINISHED_LOCK} ]]; then
  curr_environment=$(echo -e "
  ## [BEGIN] RES Environment VDI Configuration - Do Not Delete
  IDEA_MODULE_NAME={{ context.module_name }}
  IDEA_MODULE_ID={{ context.module_id }}
  IDEA_MODULE_SET={{ context.module_set }}
  IDEA_MODULE_VERSION={{ context.module_version }}
  IDEA_CLUSTER_HOME={{ context.cluster_home_dir }}
  IDEA_APP_DEPLOY_DIR={{ context.app_deploy_dir }}
  ")

  {% if context.https_proxy != '' %}
  curr_environment+=$(echo -e "
  IDEA_HTTPS_PROXY={{ context.https_proxy }}
  IDEA_NO_PROXY={{ context.no_proxy }}
  ")
  {% endif %}

  curr_environment+=$(echo -e "
  ## [END] RES Environment VDI Configuration - Do Not Delete

  ## [BEGIN] RES VDI Session - Do Not Delete
  IDEA_SESSION_ID="{{ context.vars.idea_session_id }}"
  IDEA_SESSION_OWNER="{{ context.vars.session_owner }}"
  ## [END] RES VDI Session
  ")

  curr_environment+=$(echo -e "
  ## [BEGIN] COGNITO VALUES - Do Not Delete
  COGNITO_MIN_ID="{{ context.vars.cognito_min_id }}"
  COGNITO_MAX_ID="{{ context.vars.cognito_max_id }}"
  COGNITO_UID_ATTRIBUTE="{{ context.vars.cognito_uid_attribute }}"
  COGNITO_DEFAULT_USER_GROUP="{{ context.vars.cognito_default_user_group }}"
  USER_POOL_ID="{{ context.config.get_string('identity-provider.cognito.user_pool_id', required=True) }}"
  VDI_CLIENT_ID="{{ context.config.get_string('identity-provider.cognito.vdi_client_id', required=True) }}"
  ## [END] COGNITO VALUES
  ")

  # Merge Environments
  /bin/bash  "${SCRIPT_DIR}/../common/merge_environments.sh" -r "${curr_environment}" -o /etc/environment

  source /etc/environment

  if [[ -f /etc/profile.d/proxy.sh ]]; then
      source /etc/profile.d/proxy.sh
  fi

  timestamp=$(date +%s)
  exec > ${BOOTSTRAP_DIR}/logs/configure.log.${timestamp} 2>&1

  echo -n "no" > ${BOOTSTRAP_DIR}/reboot_required.txt

  source "${SCRIPT_DIR}/../common/bootstrap_common.sh"

  # Begin: Configure AWS Systems Manager Agent
  /bin/bash "${SCRIPT_DIR}/../common/configure_amazon_ssm_agent.sh" -o "${RES_BASE_OS}" -s "${SCRIPT_DIR}"
  # End: Configure AWS Systems Manager Agent

  # Begin: Enable executable on vdi-helper python scripts
  chmod +x ${SCRIPT_DIR}/../vdi-helper/custom_credential_broker.py
  chmod +x ${SCRIPT_DIR}/../vdi-helper/vdi_auto_stop.py
  # End: Enable executable on vdi-helper python scripts

  # Begin: Configure CloudWatch Agent
  {%- set cloudwatch_agent_config = context.get_cloudwatch_agent_config(additional_log_files=additional_log_files) %}
  {%- if cloudwatch_agent_config %}
  echo '{{ context.utils.to_json(cloudwatch_agent_config, indent=True) }}' > /opt/aws/amazon-cloudwatch-agent/etc/amazon-cloudwatch-agent.json
  /opt/aws/amazon-cloudwatch-agent/bin/amazon-cloudwatch-agent-ctl -a fetch-config -m ec2 -s -c file:/opt/aws/amazon-cloudwatch-agent/etc/amazon-cloudwatch-agent.json
  {%- else %}
  log_warning "cloudwatch_agent_config not provided."
  {%- endif %}
  # End: Configure CloudWatch Agent

  {% include '_templates/linux/idea_proxy.jinja2' %}

  {% include '_templates/linux/idea_service_account.jinja2' %}

  {% include '_templates/linux/mount_shared_storage.jinja2' %}

  {%- include '_templates/linux/restrict_ssh_access_to_session_owner.jinja2' %}

  {%- include '_templates/linux/disable_ssh.jinja2' %}

  {%- with  ebs_volume_tags = [
    {'Key':'res:EnvironmentName', 'Value': context.cluster_name },
    {'Key':'res:ModuleName', 'Value': context.module_name },
    {'Key':'res:ModuleId', 'Value': context.module_id },
    {'Key':'res:Project', 'Value': context.vars.project },
    {'Key':'Name', 'Value': context.cluster_name + '/' + context.module_id + ' Root Volume' }
  ] %}
    {% include '_templates/linux/tag_ebs_volumes.jinja2' %}
  {%- endwith %}

  {%- with  network_interface_tags = [
    {'Key':'res:EnvironmentName', 'Value': context.cluster_name },
    {'Key':'res:ModuleName', 'Value': context.module_name },
    {'Key':'res:ModuleId', 'Value': context.module_id },
    {'Key':'Name', 'Value': context.cluster_name + '/' + context.module_id + ' Network Interface' }
  ] %}
    {% include '_templates/linux/tag_network_interface.jinja2' %}
  {%- endwith %}

  {% include '_templates/linux/chronyd.jinja2' %}

  {% include '_templates/linux/disable_ulimit.jinja2' %}

  {% include '_templates/linux/disable_strict_host_check.jinja2' %}

  {% include '_templates/linux/disable_motd_update.jinja2' %}

  {%- with secure_path = '/bin:/usr/bin:/sbin:/usr/sbin:/usr/local/bin' %}
    {% include '_templates/linux/sudoer_secure_path.jinja2' %}
  {%- endwith %}

  {%- with messages = [
    context.module_name + ' (v'+context.module_version+'), Cluster: ' + context.cluster_name
  ] %}
    {% include '_templates/linux/motd.jinja2' %}
  {%- endwith %}

  # TODO: enable virtual-desktop-app for ubuntu after resolve dependency issue
  if [[ $RES_BASE_OS != "ubuntu2204" ]]; then
      /bin/bash ${SCRIPT_DIR}/install_app.sh "{{context.vars.app_package_uri}}"
  fi

  {% if context.config.get_string('directoryservice.name', default='') != '' %}
    {% include '_templates/linux/join_directoryservice.jinja2' %}
  {% endif %}

  {% if context.config.get_string('scheduler.provider') == 'openpbs' %}
    {% include '_templates/linux/openpbs_client.jinja2' %}
  {% endif %}

  {% if context.config.get_bool('identity-provider.cognito.enable_native_user_login') %}
    /bin/bash "${SCRIPT_DIR}/../common/configure_cognito_modules.sh" -o "${RES_BASE_OS}" -s "${SCRIPT_DIR}" -u "${IDEA_SESSION_OWNER}" -r "${AWS_REGION}" -i "${COGNITO_MIN_ID}" -a "${COGNITO_MAX_ID}" -d "${COGNITO_UID_ATTRIBUTE}" -p "${USER_POOL_ID}" -c "${VDI_CLIENT_ID}" -g "${COGNITO_DEFAULT_USER_GROUP}" -x "${https_proxy}"
  {% endif %}

  {% include '_templates/linux/set_sudoers.jinja2' %}

  #Begin: Configure SSH Key Gen
  /bin/bash "${SCRIPT_DIR}/../common/configure_ssh_keygen.sh" -o "${RES_BASE_OS}" -s "${SCRIPT_DIR}" -u "${IDEA_SESSION_OWNER}"
  #End: Configure SSH Key Gen

  BROKER_CERTIFICATE_LOCATION_LOCAL='/etc/dcv/dcv_broker/dcvsmbroker_ca.pem'
  BROKER_HOSTNAME="{{ context.config.get_cluster_internal_endpoint().replace('https://', '') }}"
  INTERNAL_ALB_ENDPOINT="{{ context.config.get_cluster_internal_endpoint() }}"
  BROKER_AGENT_CONNECTION_PORT="{{ context.config.get_int('virtual-desktop-controller.dcv_broker.agent_communication_port', required=True) }}"

  function download_broker_certificate() {
    # local BROKER_CERTIFICATE_URL="{{ context.vars.broker_cert_url }}"
    # if [[ "${BROKER_CERTIFICATE_URL}" == s3://* ]]; then
    #  local AWS=$(command -v aws)
    #  $AWS s3 cp "${BROKER_CERTIFICATE_URL}" "${BROKER_CERTIFICATE_LOCATION_LOCAL}"
    # else
    #  cp "${BROKER_CERTIFICATE_URL}" "${BROKER_CERTIFICATE_LOCATION_LOCAL}"
    # fi
    echo "Downloading broker certificate logic goes here.. NO-OP"
  }

  function configure_dcv_host() {
    local HOSTNAME=$(hostname -s)
    local IDLE_TIMEOUT="{{ context.config.get_string('virtual-desktop-controller.dcv_session.idle_timeout', required=True) }}"
    local IDLE_TIMEOUT_WARNING="{{ context.config.get_string('virtual-desktop-controller.dcv_session.idle_timeout_warning', required=True) }}"
    echo "# configuring dcv host ..."

    if [[ -f /etc/dcv/dcv.conf ]]; then
      mv /etc/dcv/dcv.conf /etc/dcv/dcv.conf.${timestamp}
    fi

    local DCV_STORAGE_ROOT="/home/${IDEA_SESSION_OWNER}/storage-root"
    if [[ -L ${DCV_STORAGE_ROOT} ]]; then
      echo "something fishy is going on here. a sym-link should not exist. check with the session owner for bad actor or unwarranted usage of system."
      exit 1
    fi
    if [[ ! -d ${DCV_STORAGE_ROOT} ]]; then
      mkdir -p ${DCV_STORAGE_ROOT} # Create the storage root location if needed
      chown ${IDEA_SESSION_OWNER} ${DCV_STORAGE_ROOT}
    fi
    local GL_DISPLAYS_VALUE=":0.0"
    echo -e "[license]
  [log]
  level=debug
  [session-management]
  virtual-session-xdcv-args=\"-listen tcp\"
  [session-management/defaults]
  [session-management/automatic-console-session]
  [display]
  # add more if using an instance with more GPU
  cuda-devices=[\"0\"]
  [display/linux]
  gl-displays = [\"${GL_DISPLAYS_VALUE}\"]
  [display/linux]
  use-glx-fallback-provider=true
  [connectivity]
  enable-quic-frontend=true
  idle-timeout=${IDLE_TIMEOUT}
  idle-timeout-warning=${IDLE_TIMEOUT_WARNING}
  [security]
  supervision-control=\"enforced\"
  # ca-file=\"${BROKER_CERTIFICATE_LOCATION_LOCAL}\"
  auth-token-verifier=\"${INTERNAL_ALB_ENDPOINT}:${BROKER_AGENT_CONNECTION_PORT}/agent/validate-authentication-token\"
  no-tls-strict=true
  os-auto-lock=false
  administrators=[\"dcvsmagent\"]
  [windows]
  disable-display-sleep=true
  " > /etc/dcv/dcv.conf
  }

  function configure_dcv_agent() {
    if [[ -f "/etc/dcv-session-manager-agent/agent.conf" ]]; then
      mv /etc/dcv-session-manager-agent/agent.conf /etc/dcv-session-manager-agent/agent.conf.${timestamp}
    fi

    echo -e "version = '0.1'
  [agent]
  # hostname or IP of the broker. This parameter is mandatory.
  broker_host = '${BROKER_HOSTNAME}'
  # The port of the broker. Default: 8445
  broker_port = ${BROKER_AGENT_CONNECTION_PORT}
  # CA used to validate the certificate of the broker.
  # ca_file = '${BROKER_CERTIFICATE_LOCATION_LOCAL}'
  tls_strict = false
  # Folder on the file system from which the tag files are read.
  tags_folder = '/etc/dcv-session-manager-agent/tags/'
  broker_update_interval = 15
  [log]
  level = 'debug'
  rotation = 'daily'
  " > /etc/dcv-session-manager-agent/agent.conf

    mkdir -p /etc/dcv-session-manager-agent/tags/
    if [[ -f "/etc/dcv-session-manager-agent/tags/idea_tags.toml" ]]; then
      mkdir -p /etc/dcv-session-manager-agent/archive-tags/
      mv /etc/dcv-session-manager-agent/tags/idea_tags.toml /etc/dcv-session-manager-agent/archive-tags/idea_tags.toml.${timestamp}
    fi

    echo -e "idea_session_id=\"${IDEA_SESSION_ID}\"" > /etc/dcv-session-manager-agent/tags/idea_tags.toml
  }

  function configure_usb_remotization() {
    echo "Searching for USB Remotization configurations..."
    {%- for usb_info in context.config.get_list('vdc.server.usb_remotization', default=[]) %}
      echo -en "{{ usb_info }}\n" >>/etc/dcv/usb-devices.conf
    {%- endfor %}

  }

  function configure_x_server() {
    # Configure the X server to start automatically when the Linux server boots
    sudo systemctl set-default graphical.target
  }

  function start_x_server() {
    local output=$(x_server_validation_command)
    if [[ ! "$output" == "SI:localuser:dcv" ]]; then
      echo "# start x server ..."
      sudo systemctl isolate graphical.target
      echo "# wait for x server to start ..."
    fi
  }

  function restart_x_server() {
    echo "# restart x server ..."
    local NVIDIAXCONFIG=$(which nvidia-xconfig)
    sudo $NVIDIAXCONFIG --preserve-busid --enable-all-gpus
    sudo systemctl isolate multi-user.target
    sudo systemctl isolate graphical.target
    echo "# wait for x server to start ..."
  }

  function start_and_validate_x_server() {
    start_x_server
    verify_x_server_is_up
  }

  function start_and_configure_dcv_service() {
    echo "# start dcv server ..."
    sudo systemctl enable dcvserver
    echo """
  #  This file is part of systemd.
  #  AGENT
  #  systemd is free software; you can redistribute it and/or modify it
  #  under the terms of the GNU Lesser General Public License as published by
  #  the Free Software Foundation; either version 2.1 of the License, or
  #  (at your option) any later version.

  [Unit]
  Description=NICE DCV server daemon
  DefaultDependencies=no
  Conflicts=umount.target
  After=network-online.target remote-fs.target
  Before=umount.target

  [Service]
  PermissionsStartOnly=true
  ExecStartPre=-/sbin/modprobe eveusb
  ExecStart=/usr/bin/dcvserver -d --service
  Restart=always
  BusName=com.nicesoftware.DcvServer
  User=dcv
  Group=dcv

  [Install]
  WantedBy=multi-user.target
  """ > /usr/lib/systemd/system/dcvserver.service

    # Because we have modified the .service file we need to tell systemctl to reload and recreate the dependency tree again.
    # This is necessary because we are introducing a dependency on network.targets.
    # Refer - https://serverfault.com/questions/700862/do-systemd-unit-files-have-to-be-reloaded-when-modified
    sudo systemctl daemon-reload
    sudo systemctl restart dcvserver
  }

  function start_and_configure_dcv_agent_service() {
    echo "# start dcv session manager agent ..."
    sudo systemctl enable dcv-session-manager-agent
    echo """
  #  This file is part of systemd.
  #  AGENT
  #  systemd is free software; you can redistribute it and/or modify it
  #  under the terms of the GNU Lesser General Public License as published by
  #  the Free Software Foundation; either version 2.1 of the License, or
  #  (at your option) any later version.

  [Unit]
  Description=Agent component of DCV Session Manager
  DefaultDependencies=no
  Conflicts=umount.target
  After=network-online.target remote-fs.target
  Before=umount.target

  [Service]
  Type=simple
  ExecStart=/usr/libexec/dcv-session-manager-agent/dcvsessionmanageragent
  User=dcvsmagent
  Group=dcvsmagent

  [Install]
  WantedBy=multi-user.target
  """ > /usr/lib/systemd/system/dcv-session-manager-agent.service

    # Because we have modified the .service file we need to tell systemctl to reload and recreate the dependency tree again.
    # This is necessary because we are introducing a dependency on network.targets.
    # Refer - https://serverfault.com/questions/700862/do-systemd-unit-files-have-to-be-reloaded-when-modified
    sudo systemctl daemon-reload
    sudo systemctl restart dcv-session-manager-agent
  }

  function x_server_validation_command() {
    echo $(DISPLAY=:0 XAUTHORITY=$(ps aux | grep "X.*\-auth" | grep -v grep | sed -n 's/.*-auth \([^ ]\+\).*/\1/p') xhost | grep "SI:localuser:dcv$")
  }

  function verify_x_server_is_up() {
    max_timeout_in_seconds=120
    start_time=$(date +%s)
    echo "# validating if x server is running ..."
    sleep 10
    local output=$(x_server_validation_command)
    local count=0
    while [[ ! "$output" == "SI:localuser:dcv" ]]
    do
      echo -ne "Waiting for X Server to come up.. sleeping for 10 more seconds; ${count} seconds already slept \033[0K\r"
      count=$(($count+10))
      sleep 10
      output=$(x_server_validation_command)
      if [[ $(expr $count % 50) == 0 ]]
      then
        echo "Waited 5 times in a row. Was unsuccessful. trying to restart x server again..."
        restart_x_server
      fi
      current_time=$(date +%s)
      elapsed=$((current_time - start_time))
      if [[ $elapsed -ge $max_timeout_in_seconds ]]
      then
          echo "Max timeout for verify server reached after $elapsed seconds"
          break
      fi
    done
    current_time=$(date +%s)
    elapsed=$((current_time - start_time))
    if [[ $elapsed -lt $max_timeout_in_seconds ]]
    then
        echo "x server is up and running...."
    fi
  }

  ## -- DCV RELATED EXECUTION BEGINS HERE -- ##

  /bin/bash "${SCRIPT_DIR}/../nice-dcv-linux/disable_wayland_protocol.sh" -o "${RES_BASE_OS}" -s "${SCRIPT_DIR}"

  download_broker_certificate
  configure_dcv_host
  configure_dcv_agent
  configure_usb_remotization
  {% if context.is_gpu_instance_type() -%}
    /bin/bash "${SCRIPT_DIR}/../nice-dcv-linux/configure_gl.sh" -o "${RES_BASE_OS}" -s "${SCRIPT_DIR}"
  {% else %}
    echo "OS is "${RES_BASE_OS}", BUT not a GL Machine. No need for this configuration. NO-OP..."
  {% endif %}
  machine=$(uname -m) #x86_64 or aarch64
  if [[ $machine == "x86_64" ]]; then

    configure_x_server

    start_and_validate_x_server
  fi
  start_and_configure_dcv_service
  start_and_configure_dcv_agent_service

  ## -- DCV RELATED EXECUTION ENDS HERE -- ##

  # run user customizations if available
  if [[ -f ${IDEA_CLUSTER_HOME}/dcv_host/userdata_customizations.sh ]]; then
    /bin/bash ${IDEA_CLUSTER_HOME}/dcv_host/userdata_customizations.sh >> ${BOOTSTRAP_DIR}/logs/userdata_customizations.log 2>&1
  fi

  ## run custom user configure scripts
  source /etc/launch_script_environment
  /bin/bash ${SCRIPT_DIR}/${ON_VDI_CONFIGURED_COMMANDS}

  # notify controller
  CONTROLLER_EVENTS_QUEUE_URL="{{ context.config.get_string('virtual-desktop-controller.events_sqs_queue_url', required=True) }}"
  MESSAGE="{{ context.vars.dcv_host_ready_message }}"
  AWS=$(command -v aws)
  $AWS sqs send-message --queue-url ${CONTROLLER_EVENTS_QUEUE_URL} --message-body ${MESSAGE} --region ${AWS_REGION} --message-group-id ${IDEA_SESSION_ID}

  mkdir -p ${SEMAPHORE_DIR}
  echo $(date +%s) > ${CONFIG_FINISHED_LOCK}

  REBOOT_REQUIRED=$(cat /root/bootstrap/reboot_required.txt)
  if [[ "${REBOOT_REQUIRED}" == "yes" ]]; then
    log_info "Triggering 3rd Reboot Required for RES Configuration for VDI"
    echo -n "no" > /root/bootstrap/reboot_required.txt

    reboot

    # Exit to avoid running the rest of script before reboot is completed.
    exit 0
  fi

else
  # Rerun script if RERUN_ON_REBOOT is true when reboot
  source /etc/launch_script_environment
  if [[ "${RERUN_ON_REBOOT}" == "True" && -f ${INSTANCE_READY_LOCK} ]]; then
    echo "Rerunning ON_VDI_CONFIGURED_COMMANDS"
    /bin/bash ${SCRIPT_DIR}/${ON_VDI_CONFIGURED_COMMANDS}
  else
    echo "Skipping running ON_VDI_CONFIGURED_COMMANDS"
  fi
fi

/bin/bash ${SCRIPT_DIR}/configure_post_reboot.sh