#!/usr/bin/env bash
set -uo pipefail

# Check if pod name is passed as an argument
if [ $# -lt 1 ]; then
  echo "Usage: $0 <pod_name> [kube_context]" >&2
  exit 1
fi

POD_NAME="$1"  # The first argument is the name of the pod
KUBE_CONTEXT="${2:-}"  # The second argument is the kube context, default is empty

# Checks if socat is installed
if ! command -v socat > /dev/null; then
  echo "Using 'port-forward' mode to run ssh session on Kubernetes instances requires 'socat' to be installed. Please install 'socat'" >&2
  exit
fi

# Checks if netcat is installed (may not be present in many docker images)
if ! command -v nc > /dev/null; then
  echo "Using 'port-forward' mode to run ssh session on Kubernetes instances requires 'nc' to be installed. Please install 'nc' (netcat)." >&2
  exit
fi

# Establishes connection between local port and the ssh jump pod using kube port-forward
# Instead of specifying a port, we let kubectl select a random port.
# This is preferred because of socket re-use issues in kubectl port-forward,
# see - https://github.com/kubernetes/kubernetes/issues/74551#issuecomment-769185879
KUBECTL_OUTPUT=$(mktemp)
if [ -n "$KUBE_CONTEXT" ]; then
  kubectl --context="$KUBE_CONTEXT" port-forward pod/"${POD_NAME}" :22 > "${KUBECTL_OUTPUT}" 2>&1 &
else
  kubectl port-forward pod/"${POD_NAME}" :22 > "${KUBECTL_OUTPUT}" 2>&1 &
fi

# Capture the PID for the backgrounded kubectl command
K8S_PORT_FWD_PID=$!

# Wait until kubectl port-forward is ready
while ! grep -q "Forwarding from 127.0.0.1" "${KUBECTL_OUTPUT}"; do
    sleep 0.1
    # Handle the case where kubectl port-forward fails
    # It may fail if the kubeconfig is missing or invalid, or if the cluster is not reachable.
    if ! kill -0 "$K8S_PORT_FWD_PID" 2> /dev/null; then
        echo "kubectl port-forward failed. Error: $(cat "$KUBECTL_OUTPUT")" >&2
        rm -f "${KUBECTL_OUTPUT}"
        exit 1
    fi
done

# Extract the local port number assigned by kubectl.
local_port=$(awk -F: '/Forwarding from 127.0.0.1/{print $2}' "${KUBECTL_OUTPUT}" | awk -F' ' '{print $1}')

# Clean up the temporary file
rm -f "${KUBECTL_OUTPUT}"

# Add handler to terminate the port-forward process when this script exits.
trap "[[ -e /proc/$K8S_PORT_FWD_PID ]] && kill $K8S_PORT_FWD_PID" EXIT

# Checks if a connection to local_port of 127.0.0.1:[local_port] is established
while ! nc -z 127.0.0.1 "${local_port}"; do
    sleep 0.1
done

# Establishes two directional byte streams to handle stdin/stdout between
# terminal and the jump pod.
# socat process terminates when port-forward terminates.
socat - tcp:127.0.0.1:"${local_port}"
