#!/usr/bin/env bash
#
# ssh-copy-id-win — ssh-copy-id that also works against Windows OpenSSH
# servers (regular users and Administrators), while keeping full *nix
# compatibility.
#
# Windows design notes:
#
#   * cmd.exe only. Every remote command is dispatched as
#         cmd.exe /c "<inner>"
#     so it does not matter what sshd has configured as DefaultShell.
#
#   * ZERO quotes inside <inner>. This is deliberate and load-bearing:
#     if the remote DefaultShell is PowerShell, sshd runs
#         powershell -c "cmd.exe /c \"<inner>\""
#     and PowerShell re-escapes nested quotes as \" — which cmd.exe does
#     not understand. With exactly one quote pair (the outer one) the
#     string survives both shells identically. To achieve that:
#       - `cd /d %USERPROFILE%` first, then use short relative paths that
#         cannot contain spaces (.ssh\authorized_keys). `cd` consumes the
#         rest of its argument, so an unquoted path with spaces is fine.
#       - search on the base64 key blob only, which is whitespace-free.
#
#   * Appends with `(echo LINE)>>file`, never `echo LINE>>file && ...`.
#     In the latter, cmd lifts the redirect out of the middle and the
#     space before `&&` ends up INSIDE the echoed text, producing a
#     trailing space. That breaks any later exact-match dedup and silently
#     appends a duplicate key on every run.
#
#   * No hard-coded drive letters: %USERPROFILE% / %ProgramData% only.
#
#   * Administrators are written to
#     %ProgramData%\ssh\administrators_authorized_keys, because the stock
#     Windows sshd_config has a `Match Group administrators` block that
#     makes that file the ONLY one consulted for admin accounts. The
#     required restrictive ACL (SYSTEM + Administrators) is applied after.

set -u

PROG=$(basename "$0")
IDENTITY=""
PORT=""
SSH_OPTS=()
FORCE_MODE=""   # "" = autodetect, "admin", "user"
DRY_RUN=0

usage() {
  cat <<EOF
Usage: $PROG [-i identity_file] [-p port] [-o ssh_option] [-A|-U] [-n] [user@]hostname

  -i identity_file  Public key to install. Defaults to the most recently
                    modified ~/.ssh/id*.pub (same rule as ssh-copy-id).
  -p port           Remote SSH port.
  -o ssh_option     Extra option passed through to ssh (repeatable).
  -A                Force the Administrator path on Windows.
  -U                Force the regular-user path on Windows.
  -n                Dry run: print what would happen, change nothing.
  -h                This help.
EOF
}

while getopts ":i:p:o:AUnh" opt; do
  case "$opt" in
    i) IDENTITY=$OPTARG ;;
    p) PORT=$OPTARG ;;
    o) SSH_OPTS+=(-o "$OPTARG") ;;
    A) FORCE_MODE="admin" ;;
    U) FORCE_MODE="user" ;;
    n) DRY_RUN=1 ;;
    h) usage; exit 0 ;;
    \?) echo "$PROG: unknown option -$OPTARG" >&2; usage >&2; exit 1 ;;
    :)  echo "$PROG: option -$OPTARG requires an argument" >&2; exit 1 ;;
  esac
done
shift $((OPTIND - 1))

TARGET=${1:-}
if [ -z "$TARGET" ]; then
  echo "$PROG: missing [user@]hostname" >&2
  usage >&2
  exit 1
fi

PORT_OPTS=()
[ -n "$PORT" ] && PORT_OPTS=(-p "$PORT")

# This script makes several round trips (probe OS, detect admin, install,
# set ACL). On a password-auth host that would mean re-entering the password
# for each one, so all of them share a single multiplexed connection: you
# authenticate once and the rest ride along.
CTL_DIR=$(mktemp -d "${TMPDIR:-/tmp}/ssh-copy-id-win.XXXXXX")
CTL_PATH="$CTL_DIR/s"

cleanup() {
  ssh -o ControlPath="$CTL_PATH" -O exit "$TARGET" >/dev/null 2>&1
  rm -rf "$CTL_DIR"
}
trap cleanup EXIT INT TERM

ssh_run() {
  # LogLevel=ERROR keeps banners/warnings (e.g. the post-quantum notice) out
  # of the output we parse for result markers.
  ssh -o LogLevel=ERROR \
      -o ControlMaster=auto -o ControlPath="$CTL_PATH" -o ControlPersist=60 \
      -o PreferredAuthentications=publickey,keyboard-interactive,password \
      "${PORT_OPTS[@]}" "${SSH_OPTS[@]}" "$TARGET" "$@"
}

# ---------------------------------------------------------------- identity
if [ -z "$IDENTITY" ]; then
  IDENTITY=$(ls -dt "$HOME"/.ssh/id*.pub 2>/dev/null | grep -v -- '-cert\.pub$' | head -n 1)
  [ -n "$IDENTITY" ] || { echo "$PROG: no identity found and none given with -i" >&2; exit 1; }
else
  case "$IDENTITY" in
    *.pub) : ;;
    *) [ -f "${IDENTITY}.pub" ] && IDENTITY="${IDENTITY}.pub" ;;
  esac
fi
[ -f "$IDENTITY" ] || { echo "$PROG: identity file not found: $IDENTITY" >&2; exit 1; }

# First non-blank line, CR stripped, whitespace collapsed.
PUBKEY_LINE=$(tr -d '\r' < "$IDENTITY" | grep -v '^[[:space:]]*$' | head -n 1 \
              | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//' -e 's/[[:space:]][[:space:]]*/ /g')
[ -n "$PUBKEY_LINE" ] || { echo "$PROG: $IDENTITY contains no key" >&2; exit 1; }

echo "$PROG: using identity: $IDENTITY" >&2

# ------------------------------------------------------------ OS detection
# Deliberately explicit: a failed connection must NOT be mistaken for
# "not *nix, therefore Windows".
PROBE_ERR="$CTL_DIR/probe.err"
UNAME_OUT=$(ssh_run uname -s 2>"$PROBE_ERR")
SSH_RC=$?

if [ $SSH_RC -eq 255 ]; then
  echo "$PROG: cannot connect to $TARGET" >&2
  [ -s "$PROBE_ERR" ] && sed "s/^/$PROG: ssh: /" "$PROBE_ERR" >&2
  echo "$PROG: (if this host uses password auth, run from an interactive" >&2
  echo "$PROG:  terminal so ssh can prompt for it)" >&2
  exit 1
fi

REMOTE_OS=""
case "$UNAME_OUT" in
  *Linux*|*Darwin*|*BSD*|*SunOS*|*AIX*|*CYGWIN*|*MINGW*|*MSYS*) REMOTE_OS="nix" ;;
esac

if [ -z "$REMOTE_OS" ]; then
  # Probe for Windows rather than assuming it.
  VER_OUT=$(ssh_run 'cmd.exe /c ver' 2>/dev/null)
  case "$VER_OUT" in
    *Windows*|*windows*) REMOTE_OS="windows" ;;
  esac
fi

if [ -z "$REMOTE_OS" ]; then
  echo "$PROG: could not determine remote OS (uname gave '$UNAME_OUT')" >&2
  exit 1
fi

# =========================================================== *nix remote ==
if [ "$REMOTE_OS" = "nix" ]; then
  echo "$PROG: remote is *nix ($UNAME_OUT), using standard behavior" >&2

  if [ "$DRY_RUN" = 1 ]; then
    echo "$PROG: [dry-run] would append $IDENTITY to ~/.ssh/authorized_keys on $TARGET" >&2
    exit 0
  fi

  RESULT=$(ssh_run '
    set -e
    umask 077
    mkdir -p ~/.ssh
    tmpkey=$(mktemp ~/.ssh/.ssh-copy-id-win.XXXXXX)
    cat > "$tmpkey"
    touch ~/.ssh/authorized_keys
    chmod 600 ~/.ssh/authorized_keys
    if grep -qxF -f "$tmpkey" ~/.ssh/authorized_keys 2>/dev/null; then
      echo SSHCPID_EXISTS
    else
      cat "$tmpkey" >> ~/.ssh/authorized_keys
      echo SSHCPID_ADDED
    fi
    rm -f "$tmpkey"
  ' < "$IDENTITY")

  case "$RESULT" in
    *SSHCPID_ADDED*)  echo "$PROG: key added." ;;
    *SSHCPID_EXISTS*) echo "$PROG: key already present, nothing to do." ;;
    *) echo "$PROG: remote command failed:" >&2; echo "$RESULT" >&2; exit 1 ;;
  esac
  exit 0
fi

# ======================================================== Windows remote ==
echo "$PROG: remote is Windows, using cmd.exe install path" >&2

# The key blob is the longest whitespace-separated field: type blob [comment].
KEYBLOB=$(printf '%s\n' $PUBKEY_LINE | awk 'length($0) > length(m) { m = $0 } END { print m }')

# Everything below is injected unquoted into a cmd.exe command line, so it
# must be free of characters cmd would treat as syntax.
case "$KEYBLOB" in
  ""|*[!A-Za-z0-9+/=]*)
    echo "$PROG: could not extract a usable base64 key blob from $IDENTITY" >&2
    exit 1 ;;
esac
case "$PUBKEY_LINE" in
  *[\"%^\&\|\<\>\(\)!]*)
    echo "$PROG: key line contains characters unsafe for an unquoted cmd.exe command line" >&2
    echo "$PROG: (one of: \" % ^ & | < > ( ) !) — edit the key's comment and retry" >&2
    exit 1 ;;
esac

# ---------------------------------------------------------- admin or user
MODE="$FORCE_MODE"
if [ -z "$MODE" ]; then
  # S-1-5-32-544 = BUILTIN\Administrators, locale-independent.
  if ssh_run 'cmd.exe /c "whoami /groups | findstr /L S-1-5-32-544 >nul"' >/dev/null 2>&1; then
    MODE="admin"
  else
    MODE="user"
  fi
fi

if [ "$MODE" = "admin" ]; then
  BASE_DIR='%ProgramData%'
  SUB_DIR='ssh'
  REL_FILE='ssh\administrators_authorized_keys'
  echo "$PROG: installing as Administrator key -> %ProgramData%\\$REL_FILE" >&2
else
  BASE_DIR='%USERPROFILE%'
  SUB_DIR='.ssh'
  REL_FILE='.ssh\authorized_keys'
  echo "$PROG: installing as regular user key -> %USERPROFILE%\\$REL_FILE" >&2
fi

# Single quote pair, none inside. See header comment.
# The `if` MUST be parenthesised. Without parens, cmd swallows everything
# after it as the if-body, so when the directory already exists the entire
# rest of the chain is skipped and the command silently does nothing.
INNER="cd /d $BASE_DIR & (if not exist $SUB_DIR mkdir $SUB_DIR) & findstr /L /C:$KEYBLOB $REL_FILE >nul 2>&1 && echo SSHCPID_EXISTS || ( (echo $PUBKEY_LINE)>>$REL_FILE & echo SSHCPID_ADDED )"

if [ "$DRY_RUN" = 1 ]; then
  echo "$PROG: [dry-run] mode=$MODE, would run on $TARGET:" >&2
  echo "  cmd.exe /c \"$INNER\"" >&2
  exit 0
fi

WIN_ERR=$(mktemp)
RESULT=$(ssh_run "cmd.exe /c \"$INNER\"" 2>"$WIN_ERR")
case "$RESULT" in
  *SSHCPID_ADDED*)  echo "$PROG: key added." ;;
  *SSHCPID_EXISTS*) echo "$PROG: key already present, nothing to do." ;;
  *)
    echo "$PROG: remote command failed." >&2
    echo "$PROG: command was:" >&2
    echo "  cmd.exe /c \"$INNER\"" >&2
    [ -s "$WIN_ERR" ] && { echo "$PROG: stderr:" >&2; cat "$WIN_ERR" >&2; }
    [ -n "$RESULT" ] && { echo "$PROG: stdout:" >&2; echo "$RESULT" >&2; }
    rm -f "$WIN_ERR"
    exit 1 ;;
esac
rm -f "$WIN_ERR"

if [ "$MODE" = "admin" ]; then
  # Windows sshd ignores administrators_authorized_keys unless only SYSTEM
  # and Administrators have access. SIDs keep this locale-independent.
  ACL_INNER="cd /d $BASE_DIR & icacls $REL_FILE /inheritance:r /grant *S-1-5-32-544:F /grant *S-1-5-18:F"
  if ACL_OUT=$(ssh_run "cmd.exe /c \"$ACL_INNER\"" 2>&1); then
    echo "$PROG: ACL set (SYSTEM + Administrators only)."
  else
    echo "$PROG: WARNING: could not set ACL on $REL_FILE — Windows will refuse the key until fixed:" >&2
    echo "$ACL_OUT" >&2
    exit 1
  fi
fi