kubessh.sh

kubessh() {
  local namespace=""
  local pod=""
  local init_container=""
  local main_container=""
  
  while getopts ":n:" opt; do
    case $opt in
      n)
        namespace="$OPTARG"
        ;;
      \?)
        echo "Invalid option: -$OPTARG" >&2
        return 1
        ;;
    esac
  done

  shift $((OPTIND-1))

  if [ -z "$1" ]; then
    echo "Usage: kubessh [-n namespace] pod-name [init-container]"
    return 1
  fi

  pod="$1"
  init_container="$2"

  if [ -z "$namespace" ]; then
    namespace="$(kubectl config view --minify --output 'jsonpath={..namespace}' 2>/dev/null)"
  fi

  if ! kubectl get pod -n "$namespace" "$pod" &>/dev/null; then
    echo "Pod '$pod' not found in namespace '$namespace'."
    return 1
  fi

  if [ -n "$init_container" ]; then
    if ! kubectl get pod -n "$namespace" "$pod" -o jsonpath='{.spec.initContainers[*].name}' | grep -q -w "$init_container"; then
      echo "Init container '$init_container' not found in pod '$pod'."
      return 1
    fi
  else
    main_container=$(kubectl get pod -n "$namespace" "$pod" -o=jsonpath='{.spec.containers[0].name}')
  fi

  if [ -z "$init_container" ]; then
    echo "No init container specified. Connecting to the main container '$main_container'."
    init_container="$main_container"
  fi

  # Execute kubectl exec with the specified container
  kubectl exec -it -n "$namespace" "$pod" -c "$init_container" -- /bin/bash
}

# Define autocomplete function for kubessh
_kubessh_autocomplete() {
  local namespaces
  namespaces=$(kubectl get namespaces -o custom-columns=:metadata.name)
  local pods
  local containers

  local current_word="${COMP_WORDS[COMP_CWORD]}"
  local prev_word="${COMP_WORDS[COMP_CWORD-1]}"

  local namespace_flag=false
  local specified_namespace=""

  for ((i=0; i<${#COMP_WORDS[@]}; i++)); do
    if [ "${COMP_WORDS[i]}" == "-n" ]; then
      # If -n is found in the command, set the namespace flag and capture the specified namespace if present
      namespace_flag=true
      if ((i + 1 < ${#COMP_WORDS[@]})); then
        specified_namespace="${COMP_WORDS[i+1]}"
      fi
      break
    fi
  done

  if [ -n "$specified_namespace" ]; then
    current_namespace="$specified_namespace"
  else
    current_namespace=$(kubectl config view --minify --output 'jsonpath={..namespace}' 2>/dev/null)
  fi

  if [ -n "$current_namespace" ]; then
    pods=$(kubectl get pods -n "$current_namespace" -o custom-columns=:metadata.name)
    containers=$(kubectl get pod -n "$current_namespace" "${COMP_WORDS[2]}" -o=jsonpath='{.spec.initContainers[*].name}')
  else
    pods=$(kubectl get pods -o custom-columns=:metadata.name)
    containers=$(kubectl get pod "${COMP_WORDS[2]}" -o=jsonpath='{.spec.initContainers[*].name}')
  fi

  COMPREPLY=($(compgen -W "$pods" -- "$current_word"))
  COMPREPLY+=($(compgen -W "$containers" -- "$current_word"))
}

# Register the autocomplete function for kubessh
complete -F _kubessh_autocomplete kubessh