Ensure Admission Control Plugin NodeRestriction Is Set

#!/usr/bin/env bash
#
# Automation: Ensure kube-apiserver enables the NodeRestriction admission plugin
# Scope: every control plane node
#
# Run this script as root on each control plane node.
# It is safe to re-run; changes are idempotent.
#
# Operational impact:
#   - Editing /etc/kubernetes/manifests/kube-apiserver.yaml will cause the
#     kube-apiserver static pod to be restarted by the kubelet.

set -euo pipefail

MANIFEST="/etc/kubernetes/manifests/kube-apiserver.yaml"
BACKUP_DIR="/etc/kubernetes/manifests/backup-node-restriction"
TIMESTAMP="$(date +%Y%m%d-%H%M%S)"

echo "[*] Ensuring NodeRestriction admission plugin is enabled in ${MANIFEST}"

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

if [[ ! -f "${MANIFEST}" ]]; then
  echo "ERROR: kube-apiserver manifest not found at ${MANIFEST}" >&2
  exit 1
fi

mkdir -p "${BACKUP_DIR}"

backup_file="${BACKUP_DIR}/kube-apiserver.yaml.${TIMESTAMP}"
cp -p "${MANIFEST}" "${backup_file}"
echo "[*] Backed up current manifest to ${backup_file}"

# Function to check if NodeRestriction is already present in an --enable-admission-plugins arg line
has_noderestriction_enabled() {
  grep -E '^\s*- --enable-admission-plugins=' "${MANIFEST}" | grep -q 'NodeRestriction'
}

# Function to check if there's any --enable-admission-plugins arg line
has_enable_plugins_arg() {
  grep -qE '^\s*- --enable-admission-plugins=' "${MANIFEST}"
}

if has_noderestriction_enabled; then
  echo "[*] NodeRestriction already present in --enable-admission-plugins; no manifest change needed."
else
  echo "[*] Updating ${MANIFEST} to ensure NodeRestriction is enabled..."

  tmp_manifest="$(mktemp)"
  trap 'rm -f "${tmp_manifest}"' EXIT

  if has_enable_plugins_arg; then
    # Append NodeRestriction to existing comma-separated list if not present
    awk '
      /^\s*- --enable-admission-plugins=/ {
        line=$0
        # remove leading list marker and spaces to inspect value
        sub(/^\s*-\s*--enable-admission-plugins=/, "", line)
        # if NodeRestriction not already there, append it
        if (line !~ /(^|,)NodeRestriction(,|$)/) {
          sub(/\r$/, "", line)
          $0 = gensub(/(--enable-admission-plugins=)(.*)/,
                     "\\1" line ",NodeRestriction", 1)
        }
      }
      { print }
    ' "${MANIFEST}" > "${tmp_manifest}"
  else
    # No --enable-admission-plugins line: insert a new one under the "command:" section
    awk '
      /command:/ && in_container == 0 {
        in_container = 1
        print
        inserted = 0
        next
      }
      in_container == 1 && /^\s*-/ && inserted == 0 {
        # Insert the new argument before the first existing command arg
        print "    - --enable-admission-plugins=NodeRestriction"
        inserted = 1
      }
      { print }
    ' "${MANIFEST}" > "${tmp_manifest}"

    # Fallback: if we failed to insert (no command: section matched), append at the end
    if ! grep -q -- "--enable-admission-plugins=NodeRestriction" "${tmp_manifest}"; then
      echo "    - --enable-admission-plugins=NodeRestriction" >> "${tmp_manifest}"
    fi
  fi

  mv "${tmp_manifest}" "${MANIFEST}"
  sync
  echo "[*] Updated ${MANIFEST}. Kubelet will restart the kube-apiserver static pod automatically."
fi

echo "[*] Waiting for kube-apiserver to be running with NodeRestriction enabled..."

# Give kubelet some time to restart the static pod if it changed
sleep 10

# Verification: check running kube-apiserver process flags
if /bin/ps -ef | grep kube-apiserver | grep -v grep | grep -q -- "--enable-admission-plugins"; then
  if /bin/ps -ef | grep kube-apiserver | grep -v grep | grep -- "--enable-admission-plugins" | grep -q "NodeRestriction"; then
    echo "[OK] kube-apiserver process has --enable-admission-plugins including NodeRestriction:"
    /bin/ps -ef | grep kube-apiserver | grep -v grep | grep -- "--enable-admission-plugins" || true
    exit 0
  else
    echo "[WARN] kube-apiserver is running but --enable-admission-plugins does not include NodeRestriction yet."
    echo "Current flag lines:"
    /bin/ps -ef | grep kube-apiserver | grep -v grep | grep -- "--enable-admission-plugins" || true
    exit 1
  fi
else
  echo "[WARN] Could not find a running kube-apiserver process with --enable-admission-plugins flag."
  /bin/ps -ef | grep kube-apiserver | grep -v grep || true
  exit 1
fi