#!/usr/bin/env bash
#
# A kubectl plugin to ssh into Kubernetes nodes using a SSH jump host Pod
#

[[ -n $DEBUG ]] && set -x -e

PLUGIN_DIR="${HOME}/.kube/kubectlssh"
PLUGIN_SSH_OPTIONS_FILE="${PLUGIN_DIR}/options"
MAX_POD_CREATION_TIME=10   # unit: second
SSH_AGENT_ENV_FILE="${PLUGIN_DIR}/sshagent-env"
SSH_AGENT_PID_FILE="${PLUGIN_DIR}/sshagent-pid"

help(){
  echo "Usage: "
  echo "  kubectl ssh-jump <dest_node> [options]"
  echo ""
  options
}

options(){
  cat <<"EOF"
Options:
  <dest_node>                     Destination node name or IP address
                                  dest_node must start from the following letters:
                                  ASCII letters 'a' through 'z' or 'A' through 'Z',
                                  the digits '0' through '9', or hyphen ('-').
                                  NOTE: Setting dest_node as 'jumphost' allows to 
                                  ssh into SSH jump Pod as 'root' user 
  -u, --user <sshuser>            SSH User name
  -i, --identity <identity_file>  Identity key file, or PEM(Privacy Enhanced Mail)
  -p, --pubkey <pub_key_file>     Public key file
  -P, --port <port>               SSH port for target node SSH server
                                  Defaults to 22
  -a, --args <args>               Args to exec in ssh session
  --pod-template <file>           Path to custom sshjump pod definition
  --skip-agent                    Skip automatically starting SSH agent and adding 
                                  SSH Identity key into the agent before SSH login
                                  (=> You need to manage SSH agent by yourself)
  --cleanup-agent                 Clearning up SSH agent at the end
                                  The agent is NOT cleaned up in case that
                                  --skip-agent option is given
  --cleanup-jump                  Clearning up sshjump pod at the end
                                  Defaults to skip cleaning up sshjump pod
  -v, --verbose                   Run ssh in verbose mode (=ssh -vvv)
  -h, --help                      Show this message

Example:
  Scenario1 - You have private & public SSH key on your side
  $ kubectl ssh-jump -u myuser -i ~/.ssh/id_rsa -p ~/.ssh/id_rsa.pub hostname

  Scenario2 - You have .pem file but you don't have public key on your side
  $ kubectl ssh-jump -u ec2-user -i ~/.ssh/mykey.pem hostname

EOF
}

read_options(){
  if [[ -f "${PLUGIN_SSH_OPTIONS_FILE}" ]]; then
    source ${PLUGIN_SSH_OPTIONS_FILE}
  fi  
}

write_options(){
  local sshuser="$1"
  local identity="$2"
  local pubkey="$3"
  local port="$4"

  cat << EOF > ${PLUGIN_SSH_OPTIONS_FILE}
sshuser=${sshuser}
identity=${identity}

pubkey=${pubkey}
port=${port}
EOF
}

get_node_list(){
  echo "List of destination node..."
  kubectl get no -o custom-columns=Hostname:.metadata.name,Internal-IP:'{.status.addresses[?(@.type=="InternalIP")].address}'
  echo ""
}

get_openssh_verion_number() {
  ssh -V 2>&1 | awk -F'[_,]' '{print $2+0}'
}

cleanup_sshjump_pod(){
  echo "Clearning up SSH Jump host (Pod)..."
  kubectl delete pod sshjump
}

check_and_start_agent(){
  local identity="$1"

  is_alive="no"
  if [ -f ${SSH_AGENT_PID_FILE} ]; then
    SSH_AGENT_PID=$(cat ${SSH_AGENT_PID_FILE})
    source ${SSH_AGENT_ENV_FILE}
    if [ ${SSH_AGENT_PID} -gt 0 ] && ps -p ${SSH_AGENT_PID} > /dev/null
    then
      echo "ssh-agent is already running"
      is_alive="yes"
    fi
  fi
  if [ "${is_alive}" = "no" ]; then
    ssh-agent > ${SSH_AGENT_ENV_FILE}
    agent_pid=$(cat ${SSH_AGENT_ENV_FILE} | grep 'echo Agent pid' |sed 's/echo Agent pid //; s/;//')
    echo "Started ssh-agent: pid=${agent_pid}"
    echo ${agent_pid} > ${SSH_AGENT_PID_FILE}
    source ${SSH_AGENT_ENV_FILE}
    ssh-add ${identity}
  fi
}

cleanup_agent(){
  echo "Killing ssh-agent..."
  ssh-agent -k
  if [ -f ${SSH_AGENT_PID_FILE} ]; then
    rm ${SSH_AGENT_PID_FILE}
  fi 
  if [ -f ${SSH_AGENT_ENV_FILE} ]; then
    rm ${SSH_AGENT_ENV_FILE}
  fi 
}

create_jump_pod(){
  local pod_template
  if [[ -n "${jump_pod_template:-}" && -e "${jump_pod_template}" ]]; then
    pod_template=$(<"${jump_pod_template}")
  fi
  if [[ -z "${pod_template}" ]]; then
    pod_template=$(cat <<EOF
apiVersion: v1
kind: Pod
metadata:
  name: sshjump
  labels:
    env: test
spec:
  containers:
  - name: sshjump
    image: corbinu/ssh-server
    ports:
    - containerPort: 22
  nodeSelector:
    "kubernetes.io/os": linux
EOF
)
  fi
  echo "Creating SSH jump host (Pod)..."
  echo "${pod_template}" | kubectl apply -f -
}

run_ssh_node(){
  local destnode="$1"
  local sshuser="$2"
  local identity="$3"
  local pubkey="$4"
  local port="$5"
  local sshargs="$6"

  # Install an SSH Server if not yet installed
  r=$(kubectl get pod sshjump 2>/dev/null | tail -1 | awk '{print $1}') # 
  if [ "${r}" != "sshjump" ];then
    create_jump_pod

    # Wait until sshjump gets ready
    c=1
    while [[ ${c} -le ${MAX_POD_CREATION_TIME} ]];
    do
      pod_status=$(kubectl get pod sshjump 2>/dev/null | tail -1 | awk '{print $3}')
      if [ "${pod_status}" = "Running" ]; then
        break
      fi
      (( c++ ))
      sleep 1
    done
  fi

  local identity_sshjump=${identity}
  local pubkey_sshjump=${pubkey}
  if [ ! -f "${pubkey_sshjump}" ]; then
    # Generate temp private/public key to ssh to the sshjump if the pubkey isn't given
    identity_sshjump=${PLUGIN_DIR}/id_rsa_sshjump
    pubkey_sshjump=${PLUGIN_DIR}/id_rsa_sshjump.pub
    if [ ! -f "${pubkey_sshjump}" ]; then
      echo "Generating nopass SSH pri/pub key to ssh to the sshjump ..."
      ssh-keygen -t rsa -f ${identity_sshjump} -N '' > /dev/null
    fi
  fi

  # Setup portforward
  kubectl port-forward sshjump 2222:22 2>/dev/null &
  pid_port_forward=$!

  # Wait a bit for the port forwarding to get ready for connection handling for 2222
  sleep 2

  # Inject public SSH key to sshjump
  cat ${pubkey_sshjump} | \
    kubectl exec -i sshjump -- /bin/bash -c "cat > /root/.ssh/authorized_keys"

  # Add default ssh option
  sshargs="${sshargs} -o StrictHostKeyChecking=no"
 
  # Add RSA workaround options if the local OpenSSH version >= 8.5
  sshversion=$(get_openssh_verion_number)
  if [ $(echo "${sshversion} >= 8.5" | bc) -eq 1 ]; then
    sshargs="${sshargs} -o HostkeyAlgorithms=+ssh-rsa -o PubkeyAcceptedAlgorithms=+ssh-rsa"
  fi

  if [ "${destnode}" = "sshjump" ]; then
    ssh ${sshuser}@127.0.0.1 -p 2222 -i ${identity_sshjump} ${sshargs}
  else
    # Using the SSH Server as a jumphost (via port-forward proxy), ssh into the desired Node
    ssh -i ${identity} -p ${port} ${sshuser}@${destnode} \
      -o "ProxyCommand ssh root@127.0.0.1 -p 2222 -i ${identity_sshjump} ${sshargs} \"nc %h %p\"" ${sshargs}
  fi
  # Stop port-forward
  kill -3 ${pid_port_forward} 2>/dev/null
}

plugin_main() {
  skip_agent=no
  cleanup_jump=no
  cleanup_agent=no
  sshargs=""
  while [ $# -gt 0 ] ; do
    nSkip=1
    case $1 in
      "-h" | "--help")
        help
        exit 0
	      ;;
      "-v" | "--verbose" )
        sshargs="${sshargs} -vvv"
        ;;
     "--cleanup-jump")
        cleanup_jump=yes
        ;;
      "--cleanup-agent")
        cleanup_agent=yes
        ;;
      "--skip-agent")
        skip_agent=yes
        ;;
      "-u" | "--user" )
        c_sshuser=$2
        nSkip=2
        ;;
      "-i" | "--identity" )
        c_identity=$2
        nSkip=2
        ;;
      "-p" | "--pubkey" )
        c_pubkey=$2
        nSkip=2
        ;;
      "-P" | "--port")
        c_port=$2
        nSkip=2
        ;;
      "-a" | "--args" )
        sshargs="${sshargs} $2"
        nSkip=2
        ;;
      "--pod-template")
         jump_pod_template="$2"
         nSkip=2
         ;;
      [0-9a-zA-Z-]*)
        destnode=$1
        ;;
      *)
        help >&2
        exit 1
        ;;
    esac
    shift $nSkip
  done

  if [[ "$(type kubectl &>/dev/null; echo $?)" -eq 1 ]]; then
    echo "Error: missing kubectl command" >&2
    echo "Please install kubectl (https://kubernetes.io/docs/tasks/tools/install-kubectl/)" >&2
    exit 1
  fi

  if [ ! -n "${destnode}" ]; then
    help >&2
    echo ""
    get_node_list
    exit 1
  fi
  if [ "${destnode}" = "sshjump" ]; then
    echo "Setting destination name as 'jumphost' allows to ssh into SSH jump Pod as 'root' user"
    c_sshuser=root
  fi

  if [ ! -d ${PLUGIN_DIR} ]; then
    mkdir -p ${PLUGIN_DIR}
  fi
  read_options

  if [ ! -n "${c_sshuser}" ]; then
    if [ ! -n "${sshuser}" ]; then
      c_sshuser="${USER}"  # default: Current executing user
    fi
    echo "using: sshuser=${sshuser}"
    c_sshuser="${sshuser}"
  fi
  if [ ! -f "${c_identity}" ]; then
    if [ ! -f "${identity}" ]; then
      echo "Error: identity file is required" >&2
      help >&2
      exit 1
    fi
    echo "using: identity=${identity}"
    c_identity="${identity}"
  fi
  if [ ! -f "${c_pubkey}" ]; then
    # From v0.4.0 pubkey file is optional to support PEM scenario
    # where you don't have public key on your side
    #if [ ! -n "${pubkey}" ]; then
    #  help >&2
    #  exit 1
    #fi
    if [ -f "${pubkey}" ]; then
      echo "using: pubkey=${pubkey}"
      c_pubkey="${pubkey}"
    fi
  fi
  if [ ! -n "${c_port}" ]; then
    if [ ! -n "${port}" ]; then
      port="22"    # default: 22
    fi
    echo "using: port=${port}"
    c_port="${port}"
  fi

  if [ "${sshargs}" != "" ]; then
    echo "using: args=${sshargs}"
  fi

  if [ "${jump_pod_template}" != "" ]; then
     echo "using: pod-template=${jump_pod_template}"
  fi

  # Caching current ssh options
  write_options "${c_sshuser}" "${c_identity}" "${c_pubkey}" "${c_port}"

  if [ "${skip_agent}" = "no" ]; then
    check_and_start_agent ${c_identity}
  fi
  # SSH Logging into desitnation node via Jump host
  run_ssh_node "${destnode}" "${c_sshuser}" "${c_identity}" "${c_pubkey}" "${c_port}" "${sshargs}"

  # Cleaning up resources if needed
  if [ "${cleanup_jump}" = "yes" ]; then
    cleanup_sshjump_pod
  fi
  if [ "${skip_agent}" = "no" ] && [ "${cleanup_agent}" = "yes" ]; then
    cleanup_agent
  fi
}

plugin_main "$@"