provision-dut-ssh.sh

#!/usr/bin/env bash
# provision-dut-ssh.sh (v2.3) — key-based SSH from test server -> SONiC DUT(s)

set -euo pipefail

KEY_DIR="/opt/dut_ssh"
KEY_PATH="$KEY_DIR/id_ed25519"
USER_NAME="admin"
DUTS_CSV=""
COMMENT="dut-automation@$(hostname)"

DO_RESTRICT=1
RESTRICT_FROM_IP=""       # single IP
RESTRICT_CIDR=""          # comma-separated CIDRs, e.g. "10.30.2.0/24,192.168.1.0/24"
DO_PERSIST=1
DO_SUDOERS=0
SUDO_CMDS="/usr/sbin/logrotate,/usr/bin/docker,/usr/bin/show,/usr/sbin/sonic-cli"

red() { printf "\033[31m%s\033[0m\n" "$*" >&2; }
grn() { printf "\033[32m%s\033[0m\n" "$*"; }
ylw() { printf "\033[33m%s\033[0m\n" "$*"; }

usage() {
  cat <<USAGE
Usage: $0 --duts <ip[,ip2,...]> [--user admin] [--key-path /opt/dut_ssh/id_ed25519]
            [--no-restrict | --restrict-from <ip>] [--restrict-cidr <cidr[,cidr2]>]
            [--no-persist] [--sudoers "<cmd1,cmd2,...>"] [--comment "text"]
Examples:
  $0 --duts 10.30.2.129 --restrict-cidr 10.30.2.0/24
  $0 --duts 10.30.2.75,10.30.2.129 --no-restrict
USAGE
  exit 1
}

# --- args ---
while [[ $# -gt 0 ]]; do
  case "$1" in
    --duts) DUTS_CSV="$2"; shift 2;;
    --user) USER_NAME="$2"; shift 2;;
    --key-path) KEY_PATH="$2"; KEY_DIR="$(dirname "$KEY_PATH")"; shift 2;;
    --comment) COMMENT="$2"; shift 2;;
    --no-restrict) DO_RESTRICT=0; shift;;
    --restrict-from) RESTRICT_FROM_IP="$2"; DO_RESTRICT=1; shift 2;;
    --restrict-cidr) RESTRICT_CIDR="$2"; DO_RESTRICT=1; shift 2;;
    --no-persist) DO_PERSIST=0; shift;;
    --sudoers) DO_SUDOERS=1; SUDO_CMDS="$2"; shift 2;;
    -h|--help) usage;;
    *) red "Unknown arg: $1"; usage;;
  esac
done
[[ -z "${DUTS_CSV}" ]] && usage
IFS=',' read -r -a DUTS <<< "$DUTS_CSV"

# --- key ensure ---
if [[ ! -f "$KEY_PATH" ]]; then
  ylw "Generating key at $KEY_PATH"
  sudo install -d -m 700 "$KEY_DIR"
  sudo chown "$USER":"$USER" "$KEY_DIR"
  ssh-keygen -t ed25519 -C "$COMMENT" -N '' -f "$KEY_PATH"
fi
chmod 600 "$KEY_PATH"
KEY_PUB="${KEY_PATH}.pub"
KEY_B64="$(awk '{print $2}' "$KEY_PUB")"

add_known_host() { ssh-keyscan -T 5 "$1" >> "$HOME/.ssh/known_hosts" 2>/dev/null || true; }

# Guess a CIDR for the interface used to reach HOST (e.g., 10.30.2.10/24)
guess_cidr_to() {
  local host="$1"
  local dev src
  dev=$(ip route get "$host" 2>/dev/null | awk '/dev/ {for(i=1;i<=NF;i++) if($i=="dev"){print $(i+1); exit}}')
  src=$(ip route get "$host" 2>/dev/null | awk '/src/ {for(i=1;i<=NF;i++) if($i=="src"){print $(i+1); exit}}')
  [[ -z "$dev" || -z "$src" ]] && { echo ""; return; }
  # pick first inet address/prefix on that dev
  ip -o -f inet addr show dev "$dev" | awk -v s="$src" '$4 ~ "/" {print $4}' | head -n1
}

for DUT in "${DUTS[@]}"; do
  ylw "=== DUT: $DUT ==="
  add_known_host "$DUT"

  # If key already works, skip copy; otherwise install pubkey (disable agent spam so password works once)
  if ssh -i "$KEY_PATH" -o IdentitiesOnly=yes -o BatchMode=yes "$USER_NAME@$DUT" true 2>/dev/null; then
    grn "[ok] key already works on $DUT"
  else
    ylw "[..] installing public key to $DUT (enter $USER_NAME password)"
    env -u SSH_AUTH_SOCK ssh-copy-id -i "$KEY_PUB" \
      -o PubkeyAuthentication=no -o PreferredAuthentications=password \
      "$USER_NAME@$DUT"
  fi

  # Build restriction prefix
  OPT_PREFIX=""
  if [[ $DO_RESTRICT -eq 1 ]]; then
    ALLOWED_LIST=()
    [[ -n "$RESTRICT_FROM_IP" ]] && ALLOWED_LIST+=("$RESTRICT_FROM_IP")
    if [[ -n "$RESTRICT_CIDR" ]]; then
      IFS=',' read -r -a CIDRS <<< "$RESTRICT_CIDR"
      ALLOWED_LIST+=("${CIDRS[@]}")
    fi
    if [[ ${#ALLOWED_LIST[@]} -eq 0 ]]; then
      # fallback to auto-guessed CIDR of the path to this DUT
      GUESS=$(guess_cidr_to "$DUT")
      [[ -n "$GUESS" ]] && ALLOWED_LIST+=("$GUESS")
    fi
    if [[ ${#ALLOWED_LIST[@]} -eq 0 ]]; then
      red "Could not determine any allowed source for from=\"...\"; use --restrict-cidr or --no-restrict"
      exit 1
    fi
    FROM_VAL=$(IFS=, ; echo "${ALLOWED_LIST[*]}")
    OPT_PREFIX="from=\"$FROM_VAL\",no-agent-forwarding,no-port-forwarding,no-X11-forwarding"
  fi

  # Normalize authorized_keys: remove any line containing this key blob, then append the normalized entry
  ssh -i "$KEY_PATH" -o IdentitiesOnly=yes -o StrictHostKeyChecking=yes \
    "$USER_NAME@$DUT" KEY_B64="$KEY_B64" OPT_PREFIX="$OPT_PREFIX" COMMENT="$COMMENT" 'bash -s' <<'EOS'
set -e
AUTH="$HOME/.ssh/authorized_keys"
mkdir -p "$HOME/.ssh"; chmod 700 "$HOME/.ssh"
touch "$AUTH"; chmod 600 "$AUTH"
TMP="$(mktemp)"
# Remove any existing line containing the key blob (works whether options are present or not)
awk -v k="$KEY_B64" 'index($0,k)==0 {print}' "$AUTH" >"$TMP" || true
if [ -n "$OPT_PREFIX" ]; then
  echo "$OPT_PREFIX ssh-ed25519 $KEY_B64 $COMMENT" >>"$TMP"
else
  echo "ssh-ed25519 $KEY_B64 $COMMENT" >>"$TMP"
fi
cat "$TMP" >"$AUTH"; rm -f "$TMP"
EOS
  grn "[ok] authorized_keys updated on $DUT"

  if [[ $DO_PERSIST -eq 1 ]]; then
    ylw "[..] enabling persistent authorized_keys via /host on $DUT"
    ssh -tt -i "$KEY_PATH" -o IdentitiesOnly=yes -o StrictHostKeyChecking=yes \
      "$USER_NAME@$DUT" 'sudo bash -s' <<EOS
set -e
install -d -m 700 /host/ssh-keys/$USER_NAME
cp -a /home/$USER_NAME/.ssh/authorized_keys /host/ssh-keys/$USER_NAME/authorized_keys
chown -R $USER_NAME:$USER_NAME /host/ssh-keys
chmod 700 /host/ssh-keys/$USER_NAME
chmod 600 /host/ssh-keys/$USER_NAME/authorized_keys
if [ -d /etc/ssh/sshd_config.d ]; then
  printf '%s\n' 'AuthorizedKeysFile .ssh/authorized_keys /host/ssh-keys/%u/authorized_keys' > /etc/ssh/sshd_config.d/10-authorized-keys.conf
else
  if ! grep -q '^AuthorizedKeysFile' /etc/ssh/sshd_config; then
    printf '%s\n' 'AuthorizedKeysFile .ssh/authorized_keys /host/ssh-keys/%u/authorized_keys' >> /etc/ssh/sshd_config
  fi
fi
systemctl reload ssh || systemctl reload sshd || true
EOS
    grn "[ok] persistence configured on $DUT"
  fi

  if [[ $DO_SUDOERS -eq 1 ]]; then
    ylw "[..] installing sudoers allowlist on $DUT"
    ssh -tt -i "$KEY_PATH" -o IdentitiesOnly=yes -o StrictHostKeyChecking=yes \
      "$USER_NAME@$DUT" 'sudo bash -s' <<EOS
set -e
FILE=/etc/sudoers.d/dut-automation
echo '$USER_NAME ALL=(root) NOPASSWD: $SUDO_CMDS' > "\$FILE"
chmod 440 "\$FILE"
visudo -cf "\$FILE"
EOS
    grn "[ok] sudoers allowlist added on $DUT"
  fi

  # Verify pure publickey (no password fallback) so failures are obvious
  if ssh -i "$KEY_PATH" -o IdentitiesOnly=yes -o PasswordAuthentication=no "$USER_NAME@$DUT" true 2>/dev/null; then
    grn "[ok] non-interactive SSH verified on $DUT (publickey only)"
  else
    red "[warn] publickey-only auth failed on $DUT. Most likely from= mismatch."
    red "      Try: --restrict-cidr 10.30.2.0/24   (or)   --no-restrict"
    exit 1
  fi
done

grn "All done."