#!/usr/bin/env bash

# Note: Since this is run as a .profile.d/ script it will be sourced, meaning that we cannot enable
# exit on error, have to use 'return' not 'exit', and returning non-zero doesn't have an effect.

heroku_exec_log_debug() {
  if [ "$HEROKU_EXEC_DEBUG" == "1" ]; then
    echo "[heroku-exec] ${1}"
  fi
}

heroku_exec_log_error() {
  if ([[ "$DYNO" != *run.* ]] || [ "$HEROKU_EXEC_DEBUG" == "1" ]); then
    echo "[heroku-exec] ERROR: ${1}"
  fi
}

heroku_exec_log_info() {
  if ([[ "$DYNO" != *run.* ]] || [ "$HEROKU_EXEC_DEBUG" == "1" ]); then
    echo "[heroku-exec] ${1}"
  fi
}

heroku_exec_open() {
  local localAddr=$1
  local localPort="1092"
  local privateKey="$HOME/.ssh/heroku_exec_rsa"

  mkdir -p $(dirname $privateKey)
  ssh-keygen -f ${privateKey} -t rsa -N '' -C '' > /dev/null 2>&1
  cat << EOF > .heroku_exec_data.json
{
  "dyno_key": "$(cat ${privateKey}.pub)",
  "dyno_ip": "${localAddr}",
  "dyno_user": "$(whoami)"
}
EOF

  # The HostKeyAlgorithms and PubkeyAcceptedKeyTypes options are required for OpenSSH 8.8+
  # until golang.org/x/crypto supports SHA-256 based RSA and heroku-exec-proxy is
  # upgraded to that version. See W-10346196.
  cat << EOF > $HOME/.ssh/sshd_config
HostKey ${privateKey}
AuthorizedKeysFile $HOME/.ssh/authorized_keys
Subsystem sftp /usr/lib/openssh/sftp-server
ClientAliveInterval 30
ClientAliveCountMax 3
PubkeyAcceptedKeyTypes +ssh-rsa
HostkeyAlgorithms +ssh-rsa
EOF

  if ssh -V 2>&1 | grep -q -e '^OpenSSH_7\.2.*$' -e '^OpenSSH_6\.6.*$'; then
    echo "UsePrivilegeSeparation no" >> $HOME/.ssh/sshd_config
  fi

  if [ -z "$(ps -C sshd -o pid=)" ]; then
    heroku_exec_log_debug "Starting sshd on localhost:${localPort}..."
    /usr/sbin/sshd -f $HOME/.ssh/sshd_config -o "Port ${localPort}"

    if [ $? -ne 0 ]; then
      heroku_exec_log_error "Could not start SSH! Heroku Exec will not be available."
    else
      (
        failures=0
        max_retries=3
        retry_wait=10
        retry_period=10
        start_time=$SECONDS
        while true; do
          iteration_start_time=$SECONDS
          tunnel=$(curl -s -X POST -d @.heroku_exec_data.json -H "Content-Type: application/json" -L ${HEROKU_EXEC_URL}/api/v1/${DYNO})

          # We support both jq and Python since:
          # - The run image for Heroku-24+ doesn't include Python.
          # - Providing a fallback to Python allows apps using the container stack to use slimmer custom images.
          # - jq only exists in the run image for Heroku-20+, but Heroku Exec has to support existing apps
          #   running older stacks too (as well as container apps that might not have JQ installed).
          if command -v python > /dev/null; then
            echo "$tunnel" | python -c 'import json,sys;obj=json.load(sys.stdin);print(obj["user"],obj["host"],obj["port"],obj["key"])' > /dev/null 2>&1
          else
            echo "$tunnel" | jq --exit-status 'has("user") and has("host") and has("port") and has("key")' > /dev/null 2>&1
          fi

          if [ $? != 0 ]; then
            heroku_exec_log_debug "error at=create_tunnel url=${HEROKU_EXEC_URL}/api/v1/${DYNO} json=$tunnel"
          else
            heroku_exec_log_debug "at=create_tunnel json=$tunnel"

            if command -v python > /dev/null; then
              proxyUser=$(echo "$tunnel" | python -c 'import json,sys;obj=json.load(sys.stdin);print(obj["user"])')
              proxyHost=$(echo "$tunnel" | python -c 'import json,sys;obj=json.load(sys.stdin);print(obj["host"])')
              proxyPort=$(echo "$tunnel" | python -c 'import json,sys;obj=json.load(sys.stdin);print(obj["port"])')
              proxyKey=$(echo "$tunnel" | python -c 'import json,sys;obj=json.load(sys.stdin);print(obj["key"])')
            else
              proxyUser=$(echo "$tunnel" | jq --raw-output '.user')
              proxyHost=$(echo "$tunnel" | jq --raw-output '.host')
              proxyPort=$(echo "$tunnel" | jq --raw-output '.port')
              proxyKey=$(echo "$tunnel" | jq --raw-output '.key')
            fi

            echo "${proxyKey}" > $HOME/.ssh/proxy_rsa.pub
            heroku_exec_log_debug "at=authorize_pubkey fingerprint=$(ssh-keygen -lf $HOME/.ssh/proxy_rsa.pub)"
            echo "${proxyKey}" >> $HOME/.ssh/authorized_keys
            heroku_exec_log_debug "at=tunnel_starting attempts=${failures} remote_host=${proxyUser}@${proxyHost}:${proxyPort} local_port=${localPort}"
            # The HostKeyAlgorithms and PubkeyAcceptedKeyTypes options are required for OpenSSH 8.8+
            # until golang.org/x/crypto supports SHA-256 based RSA and heroku-exec-proxy is
            # upgraded to that version. See W-10346196.
            ssh -o ServerAliveInterval=30 -o ServerAliveCountMax=3 \
                -o HostKeyAlgorithms=+ssh-rsa -o PubkeyAcceptedKeyTypes=+ssh-rsa \
                -o StrictHostKeyChecking=no -i ${privateKey} \
                -p ${proxyPort} -R 0.0.0.0:0:localhost:${localPort} \
                -q -N ${proxyUser}@${proxyHost} > /dev/null 2>&1

            heroku_exec_log_debug "at=tunnel_exit status=$?"
          fi

          if [ $(($SECONDS - $iteration_start_time)) -lt 30 ]; then
            failures=$((failures+1))
            if [ $failures -gt $max_retries ]; then
              if [ $(($SECONDS - $start_time)) -lt $retry_period ]; then
                heroku_exec_log_error "Could not connect to proxy! Waiting $retry_wait seconds before retry..."
                sleep $retry_wait
                retry_wait=$((retry_wait+10))
              fi
              failures=0
              start_time=$SECONDS
            fi
          else
            sleep 1
          fi
        done
      ) &
      heroku_exec_log_info "Starting"
    fi
  else
    heroku_exec_log_debug "The sshd service is already running"
  fi
}

export_jvm_opts() {
  local ip_addr=${1}
  local jmx_port=${HEROKU_JMX_PORT:-"1098"}
  local rmi_port=${HEROKU_RMI_PORT:-"1099"}

  export HEROKU_JMX_OPTIONS="-Dcom.sun.management.jmxremote \
-Dcom.sun.management.jmxremote.port=${jmx_port} \
-Dcom.sun.management.jmxremote.rmi.port=${rmi_port} \
-Dcom.sun.management.jmxremote.ssl=false \
-Dcom.sun.management.jmxremote.authenticate=false \
-Dcom.sun.management.jmxremote.local.only=true \
-Djava.rmi.server.hostname=${ip_addr} \
-Djava.rmi.server.port=${rmi_port}"

  if [ "$HEROKU_DISABLE_JMX" != "true" ] && [ "$HEROKU_DISABLE_JMX" != "1" ]; then
    export JAVA_TOOL_OPTIONS="${JAVA_TOOL_OPTIONS} ${HEROKU_JMX_OPTIONS}"
  fi
}

main() {
  # Improve the onboarding UX when Heroku Exec is being added to a custom Docker image that is missing the
  # the requirements listed here: https://devcenter.heroku.com/articles/exec#using-with-docker
  for required_program in curl ip ssh /usr/sbin/sshd; do
    if ! command -v "${required_program}" > /dev/null; then
      heroku_exec_log_error "Couldn't find '${required_program}'. Heroku Exec won't be available."
      return 0
    fi
  done
  if ! command -v jq > /dev/null && ! command -v python > /dev/null; then
    heroku_exec_log_error "Couldn't find one of either 'jq' or 'python'. Heroku Exec won't be available."
    return 0
  fi

  local ip_addr="$(ip -4 a show eth0 | grep inet | sed -E -e 's/.*inet //g' | sed -E -e 's/\/[0-9]+.*//g')"
  export_jvm_opts ${ip_addr}
  heroku_exec_open ${ip_addr}
}

main "$@"