#!/bin/bash
#
# egress_scan.sh — detect GKE nodes that have lost public internet egress.
#
# Companion to the "Node loses egress" runbook.
#
# A node can lose outbound internet while internal traffic keeps working. Pods
# already running on it stay 1/1 Running and pass every probe, because probes
# only hit localhost. This script probes the real path from inside a pod.
#
# Requires a valid gcloud login. Tokens expire every few hours, so run this
# first if you have been away:
#   gcloud auth login
#
# Usage:
#   ./egress_scan.sh 6                       # one cluster
#   ./egress_scan.sh all | grep -v '^OK '    # whole fleet, discovered at runtime
#
# Output on stdout, one line per node:
#   OK       node reached the target host
#   FAIL     node could not reach it. Investigate.
#   UNKNOWN  probe did not run. Usually a missing python in the image, or the
#            pod died mid-probe. Not proof of a healthy node.
#
# Coverage and errors go to stderr, so a filtered pipeline cannot hide them.
# Exit codes: 0 scanned, 2 no pods found in a cluster, 3 credentials or access.
#
# This method probes from tenant pods, so a node running only system workloads
# is not covered. Expect roughly 75 to 80 percent of nodes on a normal run.
# Full coverage needs a per-node DaemonSet. See DEV-2313.
#
set -u

N=${1:?usage: egress_scan.sh <cluster-number|all>}

PROJECT=plastic-labs-prod
REGION=us-east4
TARGET_HOST=tentacruel-litellm-prod.fly.dev
TARGET_PORT=443

# Check credentials before doing anything. An expired token makes every kubectl
# call fail, which would otherwise look like a cluster full of healthy nodes.
if ! gcloud auth print-access-token >/dev/null 2>&1; then
  echo "ERROR: gcloud credentials are expired or missing. Run: gcloud auth login" >&2
  exit 3
fi

# Fleet mode. Discover the cluster list rather than hardcoding a range.
if [ "$N" = "all" ]; then
  clusters=$(gcloud container clusters list --project "$PROJECT" \
    --format='value(name)' 2>/dev/null | sed -n 's/^cluster-\([0-9]\{1,\}\)$/\1/p' | sort -n)
  if [ -z "$clusters" ]; then
    echo "ERROR: could not list clusters. Check credentials and project access." >&2
    exit 3
  fi

  COVERAGE_FILE=$(mktemp "/tmp/egress_scan_cov_XXXXXX")
  export EGRESS_SCAN_COVERAGE_FILE="$COVERAGE_FILE"
  trap 'rm -f "$COVERAGE_FILE"' EXIT

  cluster_count=$(echo "$clusters" | wc -l | tr -d ' ')
  for c in $clusters; do "$0" "$c"; done

  awk -v n="$cluster_count" '{p+=$1; t+=$2} END {
    if (t > 0)
      printf "coverage: probed %d of %d nodes (%d%%) across %d clusters. Nodes running only system pods are not reachable from here.\n", p, t, (100*p/t), n > "/dev/stderr"
  }' "$COVERAGE_FILE"
  exit 0
fi

C=gke_plastic-labs-prod_${REGION}_cluster-$N

# New clusters have no local kubeconfig entry. Fetch it rather than asking the
# operator to do it by hand. This only writes to the local kubeconfig.
if ! kubectl config get-contexts -o name 2>/dev/null | grep -qx "$C"; then
  echo "INFO cluster-$N: no kubeconfig context, fetching credentials." >&2
  if ! creds_out=$(gcloud container clusters get-credentials "cluster-$N" \
        --region "$REGION" --project "$PROJECT" 2>&1); then
    echo "SKIP cluster-$N: could not get credentials: $(echo "$creds_out" | tail -1)" >&2
    exit 0
  fi
fi

# Separate "kubectl failed" from "the answer is zero". Only the first is an error.
if ! nodes_out=$(kubectl --context "$C" get nodes --no-headers 2>&1); then
  echo "ERROR cluster-$N: kubectl failed: $(echo "$nodes_out" | head -1)" >&2
  exit 2
fi
node_total=$(echo "$nodes_out" | grep -c .)

if [ "$node_total" -eq 0 ]; then
  echo "SKIP cluster-$N: no nodes. The cluster is new or idle." >&2
  exit 0
fi

NODES_FILE=$(mktemp "/tmp/egress_scan_${N}_XXXXXX")
trap 'rm -f "$NODES_FILE"' EXIT

# One Running tenant pod per node. First match wins.
kubectl --context "$C" get pods -A -o wide --no-headers 2>/dev/null \
  | awk '$1 ~ /batch/ && $4=="Running" {print $8, $1, $2}' \
  | sort -u -k1,1 \
  | while read -r node ns pod; do
      echo "$node|$ns|$pod"
    done > "$NODES_FILE"

if [ ! -s "$NODES_FILE" ]; then
  echo "SKIP cluster-$N: $node_total nodes, no Running tenant pods to probe from." >&2
  exit 0
fi

# Record coverage so a partial scan is never mistaken for a complete one.
node_probed=$(wc -l < "$NODES_FILE" | tr -d ' ')
if [ -n "${EGRESS_SCAN_COVERAGE_FILE:-}" ]; then
  echo "$node_probed $node_total" >> "$EGRESS_SCAN_COVERAGE_FILE"
else
  echo "coverage cluster-$N: probed $node_probed of $node_total nodes." >&2
fi

probe() {
  IFS='|' read -r node ns pod <<< "$1"
  ctx=$2
  cluster_num=$3
  out=$(kubectl --request-timeout=25s --context "$ctx" -n "$ns" exec "$pod" -- python -c "
import socket
try:
    s = socket.create_connection(('${TARGET_HOST}', ${TARGET_PORT}), timeout=5)
    s.close()
    print('OK')
except Exception:
    print('FAIL')
" 2>/dev/null | tail -1)
  [ -z "$out" ] && out="UNKNOWN"
  printf '%s cluster-%s %s\n' "$out" "$cluster_num" "$node"
}
export -f probe
export TARGET_HOST TARGET_PORT

# Bounded parallelism keeps load off the API server.
xargs -P 8 -I{} bash -c 'probe "$@"' _ {} "$C" "$N" < "$NODES_FILE"