#!/bin/sh
# One-time per-machine setup for unattended Aspire access:
#
#   1. SSH config  - ensure a multiplexing Host block exists (creates one if not)
#   2. Host key    - report whether the host key still needs manual acceptance
#   3. Password    - store it locally, prompting without echo
#
#   ./bin/aspire-setup-credentials
#
# Safe to re-run: every step checks before it changes anything. The password is
# stored in plain text at ~/.config/aspire/password with mode 600; anyone able
# to read your account (including root) can read it.

set -eu

HOST="${ASPIRE_HOST:-aspire}"
CRED="${ASPIRE_CRED:-$HOME/.config/aspire/password}"
CONF="${ASPIRE_SSH_CONF:-$HOME/.config/aspire/aspire.conf}"
SSH_CONFIG="$HOME/.ssh/config"

DEFAULT_HOSTNAME="aspire2a.nus.edu.sg"

if [ ! -t 0 ]; then
    echo "aspire-setup-credentials: needs an interactive terminal" >&2
    exit 1
fi

# ask <prompt> <default> -> echoes the answer
ask() {
    printf '%s [%s]: ' "$1" "$2" >&2
    IFS= read -r reply
    if [ -z "$reply" ]; then echo "$2"; else echo "$reply"; fi
}

# ask_secret <prompt> -> echoes the answer, no terminal echo
ask_secret() {
    # Disable echo *before* printing the prompt: anything typed or pasted
    # between the prompt and the stty call would otherwise appear in cleartext.
    stty -echo
    printf '%s: ' "$1" >&2
    IFS= read -r secret
    stty echo
    echo >&2
    echo "$secret"
}

restore() { stty echo 2>/dev/null || true; }
trap 'restore; echo >&2; echo "Aborted." >&2; exit 130' INT TERM

# --- 1. SSH config ---------------------------------------------------------
echo "== 1/3  SSH config =="

configured=no
if ssh -G "$HOST" 2>/dev/null | grep -qi '^controlpath [^n]' &&
   ssh -G "$HOST" 2>/dev/null | grep -qi '^controlmaster auto'; then
    configured=yes
fi

if [ "$configured" = yes ]; then
    resolved=$(ssh -G "$HOST" 2>/dev/null | awk '/^hostname /{print $2}')
    user=$(ssh -G "$HOST" 2>/dev/null | awk '/^user /{print $2}')
    echo "Host '$HOST' already configured for multiplexing ($user@$resolved). Skipping."
else
    echo "Host '$HOST' has no multiplexing config. Let's create one."

    case "$CONF" in
        *\ *)
            echo "Path contains spaces; ssh Include cannot handle that." >&2
            echo "Set ASPIRE_SSH_CONF to a path without spaces, or add the Host block" >&2
            echo "to $SSH_CONFIG by hand." >&2
            exit 1
            ;;
    esac

    hostname=$(ask "  Login node hostname" "$DEFAULT_HOSTNAME")
    username=$(ask "  Username on the cluster" "${USER:-}")

    if [ -z "$username" ]; then
        echo "Username is required; nothing written." >&2
        exit 1
    fi

    mkdir -p "$(dirname "$CONF")"
    cat > "$CONF" <<EOF
# Managed by bin/aspire-setup-credentials. Machine-local — not tracked in
# git; each machine generates its own by re-running this script.
Host $HOST
    HostName $hostname
    User $username
    ControlMaster auto
    ControlPath ~/.ssh/cm-%C
    ControlPersist yes
    ServerAliveInterval 30
    StrictHostKeyChecking accept-new
EOF
    echo "  Wrote $CONF"

    include_line="Include $CONF"
    if [ -f "$SSH_CONFIG" ] && grep -qF "$include_line" "$SSH_CONFIG"; then
        echo "  $SSH_CONFIG already includes it."
    else
        if [ -f "$SSH_CONFIG" ]; then
            backup="$SSH_CONFIG.bak.$(date +%Y%m%d%H%M%S)"
            cp "$SSH_CONFIG" "$backup"
            echo "  Backed up $SSH_CONFIG -> $backup"

            if grep -qiE "^[[:space:]]*Host[[:space:]]+.*\b$HOST\b" "$SSH_CONFIG"; then
                echo
                echo "  NOTE: $SSH_CONFIG already has a 'Host $HOST' block."
                echo "  ssh uses the first value it finds, and the Include goes at the"
                echo "  top, so the repo file wins. Remove the old block when convenient."
                echo
            fi
        fi

        # Prepend, then write back through the original file so its inode and
        # permissions survive.
        umask 077
        tmp="$SSH_CONFIG.tmp.$$"
        printf '%s\n\n' "$include_line" > "$tmp"
        [ -f "$SSH_CONFIG" ] && cat "$SSH_CONFIG" >> "$tmp"
        touch "$SSH_CONFIG"
        cat "$tmp" > "$SSH_CONFIG"
        rm -f "$tmp"
        chmod 600 "$SSH_CONFIG"
        echo "  Added Include to $SSH_CONFIG"
    fi

    if ! ssh -G "$HOST" 2>/dev/null | grep -qi '^controlmaster auto'; then
        echo "Config written but ssh still does not see it; check $SSH_CONFIG." >&2
        exit 1
    fi
fi

# --- 2. Host key -----------------------------------------------------------
echo
echo "== 2/3  Host key =="

resolved=$(ssh -G "$HOST" 2>/dev/null | awk '/^hostname /{print $2}')
KNOWN_HOSTS="$HOME/.ssh/known_hosts"

if ssh-keygen -F "$resolved" >/dev/null 2>&1; then
    echo "Host key for $resolved is already known."
    ssh-keygen -l -F "$resolved" 2>/dev/null | grep -v '^#' | sed 's/^/  /'
else
    echo "Host key for $resolved is not known yet; fetching it..."
    if keys=$(ssh-keyscan -T 10 "$resolved" 2>/dev/null) && [ -n "$keys" ]; then
        mkdir -p "$HOME/.ssh"
        chmod 700 "$HOME/.ssh"
        umask 077
        touch "$KNOWN_HOSTS"
        printf '%s\n' "$keys" >> "$KNOWN_HOSTS"
        chmod 600 "$KNOWN_HOSTS"
        echo "  Added to $KNOWN_HOSTS:"
        printf '%s\n' "$keys" | ssh-keygen -lf - 2>/dev/null | sed 's/^/  /'
    else
        echo "  Could not reach $resolved to fetch it."
        echo "  Not fatal: the config sets StrictHostKeyChecking=accept-new, so the"
        echo "  key will be recorded automatically on the first successful connect."
    fi
fi

# --- 3. Password -----------------------------------------------------------
echo
echo "== 3/3  Password =="

if [ -e "$CRED" ]; then
    printf 'Credential file already exists at %s. Overwrite? [y/N] ' "$CRED"
    IFS= read -r reply
    case "$reply" in
        y|Y|yes|YES) ;;
        *) echo "Kept the existing password. Setup complete."; exit 0 ;;
    esac
fi

pw1=$(ask_secret "Aspire password")
pw2=$(ask_secret "Confirm password")

trap - INT TERM

if [ -z "$pw1" ]; then
    echo "Empty password; nothing written." >&2
    exit 1
fi

if [ "$pw1" != "$pw2" ]; then
    echo "Passwords do not match; nothing written." >&2
    exit 1
fi

mkdir -p "$(dirname "$CRED")"
chmod 700 "$(dirname "$CRED")"

# Create the file empty with tight permissions *before* writing the secret, so
# it is never briefly readable by others.
umask 077
: > "$CRED"
chmod 600 "$CRED"
printf '%s\n' "$pw1" > "$CRED"

echo "Wrote $CRED (mode 600)."
echo
echo "Setup complete. Next:"
echo "    ./bin/aspire-connect      # should connect without prompting"
