#!/usr/bin/env bash

trap 'kubectl delete pod $container >/dev/null 2>&1 &' 0 1 2 3 15

usage() { echo -e "Usage: kubectl ssh <options> <pod name>" && grep " .)\ #" $0; exit 0; }
[ $# -eq 0 ] && usage

while getopts ":u:c:p:n:h" arg; do
 case $arg in
   p) # Specify pod name.
     POD=${OPTARG}
     ;;
   u) # Specify user
     USERNAME=${OPTARG}
     ;;
   c) # Specify container
     CONTAINER=${OPTARG}
     ;;
   n) # Specify namespaxce
     NAMESPACE=${OPTARG}
     ;;
   h) # Display help.
     usage
     exit 0
     ;;
   -- ) # Optional command to execute. Defaults to /bin/sh
     ;;
   *)
     ;;
 esac
done

current_namespace() {
  local cur_ctx

  cur_ctx="$(current_context)" || exit_err "error getting current context"
  ns="$($KUBECTL config view -o=jsonpath="{.contexts[?(@.name==\"${cur_ctx}\")].context.namespace}")" \
     || exit_err "error getting current namespace"

  if [[ -z "${ns}" ]]; then
    echo "default"
  else
    echo "${ns}"
  fi
}

current_context() {
  $KUBECTL config current-context
}

COMMAND=$(echo $@ | grep '\-\-' | sed 's|\(.*\) -- \(.*\)|\2|g')
COMMAND="${COMMAND:-/bin/sh}"

if [ -z "$POD" ] && [ -z "$CONTAINER" ] && [ -z "$USERNAME" ] && [ -z "$NAMESPACE" ]; then
   POD="$1"
fi

USERNAME="${USERNAME:-root}"

[ -z $POD ] && echo -e "\nMissing Pod Name" && exit 1

KUBECTL=$(which kubectl)

NAMESPACE="${NAMESPACE:-$(current_namespace)}"

echo -e "\nConnecting...\nPod: ${POD}\nNamespace: ${NAMESPACE}\nUser: ${USERNAME}\nContainer: $CONTAINER\nCommand: $COMMAND\n"

# Limits concurrent ssh sessions (each session deploys a pod) to 2. It's not necessary, just a preference.
test "$(exec $KUBECTL -n "${NAMESPACE}" get po "$(whoami)-1" 2>/dev/null)" && container="$(whoami)-2" || container="$(whoami)-1"

# We want to mount the docker socket on the node of the pod we're exec'ing into.
NODENAME=$( ${KUBECTL} get pod ${POD} -o go-template='{{.spec.nodeName}}' )
NODESELECTOR='"nodeSelector": {"kubernetes.io/hostname": "'$NODENAME'"},'

# Adds toleration if the target container runs on a tainted node. Assumes no more than one taint. Change if yours have more than one or are configured differently.
TOLERATION_VALUE=$($KUBECTL -n "${NAMESPACE}" get pod ${POD} -ojsonpath='{.spec.tolerations[].value}') >/dev/null 2>&1
if [[ "$TOLERATION_VALUE" ]]; then
    TOLERATION_KEY=$($KUBECTL -n "${NAMESPACE}" get pod ${POD} -ojsonpath='{.spec.tolerations[].key}')
    TOLERATION_OPERATOR=$($KUBECTL -n "${NAMESPACE}" get pod ${POD} -ojsonpath='{.spec.tolerations[].operator}')
    TOLERATION_EFFECT=$($KUBECTL -n "${NAMESPACE}" get pod ${POD} -ojsonpath='{.spec.tolerations[].effect}')
    TOLERATIONS='"tolerations": [{"effect": "'$TOLERATION_EFFECT'","key": "'$TOLERATION_KEY'","operator": "'$TOLERATION_OPERATOR'","value": "'$TOLERATION_VALUE'"}],'
else
    TOLERATIONS=''
fi

if [[ -n ${CONTAINER} ]]; then
  DOCKER_CONTAINERID=$( eval $KUBECTL -n "${NAMESPACE}" get pod ${POD} -o go-template="'{{ range .status.containerStatuses }}{{ if eq .name \"${CONTAINER}\" }}{{ .containerID }}{{ end }}{{ end }}'" )
else
  DOCKER_CONTAINERID=$( $KUBECTL -n "${NAMESPACE}" get pod ${POD} -o go-template='{{ (index .status.containerStatuses 0).containerID }}' )
fi
CONTAINERID=${DOCKER_CONTAINERID#*//}

read -r -d '' OVERRIDES <<EOF
{
    "apiVersion": "v1",
    "spec": {
        "containers": [
            {
                "image": "docker",
                "name": "'$container'",
                "stdin": true,
                "stdinOnce": true,
                "tty": true,
                "restartPolicy": "Never",
                "args": [
                  "exec",
                  "-it",
                  "-u",
                  "${USERNAME}",
                  "${CONTAINERID}",
                  "${COMMAND}"
                ],
                "volumeMounts": [
                    {
                        "mountPath": "/var/run/docker.sock",
                        "name": "docker"
                    }
                ]
            }
        ],
        $NODESELECTOR
        $TOLERATIONS
        "volumes": [
            {
                "name": "docker",
                "hostPath": {
                    "path": "/var/run/docker.sock",
                    "type": "File"
                }
            }
        ]
    }
}
EOF

eval $KUBECTL -n "${NAMESPACE}" run -it --restart=Never --image=docker --overrides="'${OVERRIDES}'" $container