Ensure Service Account key File Argument Is Appropriate

#!/usr/bin/env bash
#
# automation_fix_service_account_key_file.sh
#
# Purpose:
#   Ensure kube-apiserver manifest has an explicit --service-account-key-file
#   argument pointing at a public key file for service accounts.
#
# Usage:
#   Run on every control plane node as root:
#     sudo bash automation_fix_service_account_key_file.sh
#
# Notes:
#   - Editing /etc/kubernetes/manifests/kube-apiserver.yaml will trigger a
#     kube-apiserver restart via kubelet (static pod).
#   - Script is idempotent and safe to re-run.

set -euo pipefail

APISERVER_MANIFEST="/etc/kubernetes/manifests/kube-apiserver.yaml"
# Path to the PUBLIC key file to be used by kube-apiserver
SERVICE_ACCOUNT_KEY_FILE="/etc/kubernetes/pki/sa.pub"

echo "[INFO] Running on host: $(hostname)"

if [[ $EUID -ne 0 ]]; then
  echo "[ERROR] This script must be run as root." >&2
  exit 1
fi

if [[ ! -f "$APISERVER_MANIFEST" ]]; then
  echo "[ERROR] Manifest not found: $APISERVER_MANIFEST" >&2
  exit 1
fi

# Ensure public key exists. If not, try to derive it from private key if present.
if [[ ! -f "$SERVICE_ACCOUNT_KEY_FILE" ]]; then
  SA_KEY_PRIV="/etc/kubernetes/pki/sa.key"
  if [[ -f "$SA_KEY_PRIV" ]]; then
    echo "[INFO] Public key $SERVICE_ACCOUNT_KEY_FILE not found, generating from $SA_KEY_PRIV"
    # Requires openssl; generate RSA public key in PEM format
    openssl rsa -in "$SA_KEY_PRIV" -pubout -out "$SERVICE_ACCOUNT_KEY_FILE"
    chmod 644 "$SERVICE_ACCOUNT_KEY_FILE"
  else
    echo "[ERROR] Neither public key $SERVICE_ACCOUNT_KEY_FILE nor private key $SA_KEY_PRIV found." >&2
    echo "[ERROR] Create or provision an appropriate service account public key file first." >&2
    exit 1
  fi
else
  echo "[INFO] Found existing service account public key: $SERVICE_ACCOUNT_KEY_FILE"
fi

BACKUP="${APISERVER_MANIFEST}.$(date +%Y%m%d%H%M%S).bak"
cp "$APISERVER_MANIFEST" "$BACKUP"
echo "[INFO] Backup created at $BACKUP"

# Normalize YAML to ensure there's a 'command:' section to work with.
# Most kubeadm-based clusters already have it; we only handle that layout here.
if ! grep -qE '^\s*-\s*--service-account-key-file=' "$APISERVER_MANIFEST"; then
  echo "[INFO] --service-account-key-file not present, adding it."
  # Insert new flag under the kube-apiserver command list.
  # We append it right after the last existing -- flag line, or after 'kube-apiserver' line.
  if grep -qE '^\s*-\s*--' "$APISERVER_MANIFEST"; then
    # Append after last existing --* flag
    awk -v flag="--service-account-key-file=${SERVICE_ACCOUNT_KEY_FILE}" '
      /^\s*-\s*--/ { last_flag_line=NR }
      { lines[NR]=$0 }
      END {
        for (i=1; i<=NR; i++) {
          print lines[i]
          if (i==last_flag_line) {
            # Maintain same indentation as existing flags
            sub(/^([[:space:]]*).*/, "&", lines[last_flag_line])
            match(lines[last_flag_line], /^([[:space:]]*).*/, m)
            indent=m[1]
            print indent "- " flag
          }
        }
      }
    ' "$APISERVER_MANIFEST" > "${APISERVER_MANIFEST}.tmp"
  else
    # Fallback: just append a new command entry with the flag near the top
    awk -v flag="--service-account-key-file=${SERVICE_ACCOUNT_KEY_FILE}" '
      !inserted && /kube-apiserver/ {
        print
        print "    - " flag
        inserted=1
        next
      }
      { print }
    ' "$APISERVER_MANIFEST" > "${APISERVER_MANIFEST}.tmp"
  fi
  mv "${APISERVER_MANIFEST}.tmp" "$APISERVER_MANIFEST"
else
  echo "[INFO] --service-account-key-file already present, ensuring it points to $SERVICE_ACCOUNT_KEY_FILE"
  # Replace existing value with desired path (idempotent)
  sed -i "s#^\(\s*-\s*--service-account-key-file=\).*#\1${SERVICE_ACCOUNT_KEY_FILE}#g" "$APISERVER_MANIFEST"
fi

echo "[INFO] kube-apiserver manifest updated. kubelet will restart the static pod automatically."

# Wait for kube-apiserver process to reflect new flag
echo "[INFO] Waiting for kube-apiserver to restart with new arguments..."
RETRIES=30
SLEEP_SEC=10
success=0

for i in $(seq 1 $RETRIES); do
  if /bin/ps -ef | grep kube-apiserver | grep -v grep | grep -q -- "--service-account-key-file=${SERVICE_ACCOUNT_KEY_FILE}"; then
    success=1
    break
  fi
  echo "[INFO] Attempt $i/$RETRIES: kube-apiserver not yet running with desired flag, sleeping ${SLEEP_SEC}s..."
  sleep "$SLEEP_SEC"
done

echo
echo "================ Verification ================"
/bin/ps -ef | grep kube-apiserver | grep -v grep || true

if [[ $success -eq 1 ]]; then
  echo "[SUCCESS] kube-apiserver is running with --service-account-key-file=${SERVICE_ACCOUNT_KEY_FILE}"
  exit 0
else
  echo "[WARNING] kube-apiserver process does not yet show --service-account-key-file=${SERVICE_ACCOUNT_KEY_FILE}."
  echo "[WARNING] Check kubelet and kube-apiserver pod status manually."
  exit 1
fi