#!/bin/sh
# fillrate-task: ssh-password-off
# Turn off SSH password login (Vast verification fails while it is on), with a
# password exception for an allowlist of source IPs, and add firewall allow
# rules for sshd from those same IPs. Follows
# https://docs.vast.ai/host/disable-ssh-password-login but reloads instead of
# restarting, so existing sessions are never dropped.
#
# The allowlist is NOT in this script (this file is public). It comes from the
# signed ticket line allow_ips=<ip>_<ip>_... (the runner passes it as
# FILLRATE_ALLOW_IPS). 1 to 8 plain IPv4 addresses, "_"-separated.
#
# Steps (dry-run does 1-3 and prints the plan; apply does all):
#   1. Refuse after the ticket expiry, or when the allowlist is missing or not
#      plain IPv4.
#   2. Preflight: refuse unless at least one admin login other than fillrate
#      (a member of sudo, admin or wheel; root only when sshd allows root
#      login) has a non-empty authorized_keys with at least one key line.
#      Password-off must never leave the host reachable only by fillrate.
#   3. Read the sshd port(s) and current passwordauthentication from sshd -T.
#   4. Back up: cp -n <file> <file>.orig for sshd_config and every
#      sshd_config.d/*.conf (never overwrites an older .orig), plus a stamped
#      copy of /etc/ssh/sshd_config* and the firewall state (ufw status,
#      iptables-save, nft list ruleset) under
#      /var/lib/fillrate/backup/ssh-password-off-<stamp>.
#   5. Set PasswordAuthentication no in sshd_config and every
#      sshd_config.d/*.conf (commented, indented and "yes" lines alike). If
#      the main file has no such line, insert one before its first Match.
#      Replace any earlier fillrate block and append, at the very end of the
#      main sshd_config (the only place a Match block is valid):
#        # BEGIN fillrate ssh-password-off
#        Match Address <allowlist>
#            PasswordAuthentication yes
#        # END fillrate ssh-password-off
#   6. sshd -t. On failure, put the stamped copies back and exit 2.
#      Then systemctl reload ssh (never restart). A socket-activated sshd with
#      an inactive service needs no reload: every new connection reads config.
#   7. Verify: sshd -T says passwordauthentication no; sshd -T -C
#      addr=<ip>,user=x,host=x says yes for every allowlisted ip and no for a
#      documentation address (203.0.113.9). On failure, restore, reload, exit 2.
#   8. Firewall, ADD-only, tcp to each sshd port from each allowlisted ip:
#      ufw active -> ufw allow ... comment fillrate-ssh-allow.
#      ufw installed but inactive -> reported, never enabled; falls through.
#      iptables -> iptables -I INPUT 1 ... -m comment fillrate-ssh-allow -j ACCEPT
#      (checked with -C first, so a rerun adds nothing).
#      nft only (no iptables) -> nft insert rule inet filter input ... only
#      when that table and chain exist; otherwise reported.
#      Never changes a policy, never adds a deny/drop/reject, never deletes a
#      rule, never touches FORWARD, DOCKER*, or Vast port rules, never saves or
#      persists rulesets (iptables rules last until reboot; reported).
#   9. Print the rollback.
#
# Never touches containers, docker, or vastai, never restarts sshd, installs
# nothing, never reboots, makes no Vast API calls. The FILLRATE_* path
# overrides below are for tests; a real run uses the defaults.

set -u

SSH_DIR=${FILLRATE_SSH_DIR:-/etc/ssh}
MAIN=$SSH_DIR/sshd_config
DROP=$SSH_DIR/sshd_config.d
BACKUP_ROOT=${FILLRATE_SSH_BACKUP_ROOT:-/var/lib/fillrate/backup}
STATE=${FILLRATE_SSH_STATE:-/var/lib/fillrate-runner/ssh-password-off}
EXPIRES=${FILLRATE_EXPIRES:-0}
MODE=${FILLRATE_MODE:-apply}
[ -n "$MODE" ] || MODE=apply
ALLOW_RAW=${FILLRATE_ALLOW_IPS:-}
PROBE_OUTSIDE=203.0.113.9
TAG=fillrate-ssh-allow
BEGIN_MARK="# BEGIN fillrate ssh-password-off"
END_MARK="# END fillrate ssh-password-off"

STAMP=$(date +%Y%m%d-%H%M%S)-$$
mkdir -p "$STATE" && chmod 700 "$STATE"
LOG="$STATE/run-$STAMP.log"
: > "$LOG"
BK="$BACKUP_ROOT/ssh-password-off-$STAMP"

say() {
  printf '%s\n' "$*"
  printf '%s %s\n' "$(date +%H:%M:%S)" "$*" >> "$LOG"
}

refuse() {
  say "refusing: $1"
  say "summary: ssh-password-off result=refused reason=$1 log=$LOG"
  exit 1
}

fail_verify() {
  say "verify failed: $1"
  restore_ssh
  say "summary: ssh-password-off result=verify-failed reason=$1 restored=yes log=$LOG"
  exit 2
}

case "$MODE" in apply|dry-run) ;; *) refuse "bad-mode-$MODE" ;; esac

# 1. Window and allowlist.
now=$(date +%s)
[ "$now" -le "$EXPIRES" ] || refuse "ticket-expired"

valid_ipv4() {
  printf '%s\n' "$1" | grep -E -q '^([0-9]{1,3}\.){3}[0-9]{1,3}$' || return 1
  old_ifs=$IFS; IFS=.
  # shellcheck disable=SC2086
  set -- $1
  IFS=$old_ifs
  for o in "$@"; do
    [ "$o" -le 255 ] || return 1
    case "$o" in 0?*) return 1 ;; esac
  done
  return 0
}

[ -n "$ALLOW_RAW" ] || refuse "no-allowlist"
ALLOW=$(printf '%s\n' "$ALLOW_RAW" | tr '_' ' ')
n=0
for ip in $ALLOW; do
  valid_ipv4 "$ip" || refuse "bad-allowlist-entry"
  n=$((n + 1))
done
[ "$n" -ge 1 ] && [ "$n" -le 8 ] || refuse "allowlist-size-$n"
ALLOW_CSV=$(printf '%s\n' "$ALLOW" | tr ' ' ',')
say "mode=$MODE allowlist_count=$n"

[ -f "$MAIN" ] || refuse "no-sshd_config"
command -v sshd >/dev/null 2>&1 || refuse "no-sshd"

# 2. Preflight: a non-fillrate admin with a usable authorized_keys.
has_keys() {
  f=$1
  [ -s "$f" ] || return 1
  grep -E -q '^[[:space:]]*([^#[:space:]].*[[:space:]])?(ssh-(ed25519|rsa|dss)|ecdsa-sha2-[a-z0-9]+|sk-[a-z0-9@.-]+)[[:space:]]+[A-Za-z0-9+/=]{20,}' "$f"
}

admins=""
for g in sudo admin wheel; do
  members=$(getent group "$g" 2>/dev/null | awk -F: '{ print $4 }' | tr ',' ' ')
  admins="$admins $members"
done
sshd_t=$(sshd -T 2>/dev/null) || refuse "sshd-T-failed"
rootlogin=$(printf '%s\n' "$sshd_t" | awk '$1 == "permitrootlogin" { print $2; exit }')
[ "$rootlogin" = "no" ] || admins="$admins root"

ok_admin=""
for u in $(printf '%s\n' $admins | sort -u); do
  [ "$u" = "fillrate" ] && continue
  home=$(getent passwd "$u" 2>/dev/null | awk -F: '{ print $6 }')
  [ -n "$home" ] || continue
  if has_keys "$home/.ssh/authorized_keys" || has_keys "$home/.ssh/authorized_keys2"; then
    ok_admin="$ok_admin $u"
  fi
done
if [ -z "$ok_admin" ]; then
  say "no admin login other than fillrate has a non-empty authorized_keys (checked sudo/admin/wheel members${rootlogin:+, root when permitrootlogin=$rootlogin})"
  refuse "no-admin-key"
fi
say "preflight: admin key login present for:$ok_admin"

# 3. Current state.
ports=$(printf '%s\n' "$sshd_t" | awk '$1 == "port" { print $2 }' | sort -u | tr '\n' ' ')
[ -n "$ports" ] || ports=22
before=$(printf '%s\n' "$sshd_t" | awk '$1 == "passwordauthentication" { print $2; exit }')
say "sshd ports: $ports"
say "before: passwordauthentication $before"
confs=$(ls "$DROP"/*.conf 2>/dev/null || true)
for f in "$MAIN" $confs; do
  hits=$(grep -n -i -E '^[[:space:]]*#?[[:space:]]*PasswordAuthentication[[:space:]]' "$f" 2>/dev/null | tr '\n' ' ')
  say "  $f: ${hits:-no PasswordAuthentication line}"
done

# Firewall detection (read-only).
fw="none"; ufw_state="absent"
if command -v ufw >/dev/null 2>&1; then
  if ufw status 2>/dev/null | head -n 1 | grep -q -i 'status: active'; then ufw_state=active; else ufw_state=inactive; fi
fi
if [ "$ufw_state" = active ]; then
  fw=ufw
elif command -v iptables >/dev/null 2>&1; then
  fw=iptables
elif command -v nft >/dev/null 2>&1; then
  if nft list chain inet filter input >/dev/null 2>&1; then fw=nft; else fw=nft-no-input-chain; fi
fi
say "firewall: ufw=$ufw_state using=$fw"
[ "$ufw_state" = inactive ] && say "note: ufw is installed but inactive; it will not be enabled"
if [ "$fw" = iptables ]; then
  pol=$(iptables -S INPUT 2>/dev/null | awk '$1 == "-P" { print $3; exit }')
  say "iptables INPUT policy: ${pol:-unknown} (allow rules only matter if something later in INPUT drops ssh)"
fi

say "plan: PasswordAuthentication no in $MAIN and $(printf '%s\n' "$confs" | grep -c . ) drop-in file(s); Match Address <allowlist> PasswordAuthentication yes at the end of $MAIN; sshd -t; systemctl reload ssh; firewall add-only via $fw for tcp port(s) $ports"

if [ "$MODE" = dry-run ]; then
  say "dry-run: no changes made"
  say "summary: ssh-password-off result=dry-run before=$before firewall=$fw ufw=$ufw_state log=$LOG"
  exit 0
fi

# 4. Backups.
mkdir -p "$BK/ssh" "$BK/firewall" && chmod 700 "$BK"
cp -n "$MAIN" "$MAIN.orig"
for f in $confs; do cp -n "$f" "$f.orig"; done
cp -p "$MAIN" "$BK/ssh/sshd_config"
if [ -d "$DROP" ]; then mkdir -p "$BK/ssh/sshd_config.d"; for f in $confs; do cp -p "$f" "$BK/ssh/sshd_config.d/"; done; fi
command -v ufw >/dev/null 2>&1 && ufw status verbose > "$BK/firewall/ufw-status.txt" 2>&1
command -v iptables-save >/dev/null 2>&1 && iptables-save > "$BK/firewall/iptables-save.txt" 2>/dev/null
command -v nft >/dev/null 2>&1 && nft list ruleset > "$BK/firewall/nft-ruleset.txt" 2>/dev/null
say "backup: $BK (and .orig next to each sshd file)"

restore_ssh() {
  cp -p "$BK/ssh/sshd_config" "$MAIN"
  for f in $confs; do cp -p "$BK/ssh/sshd_config.d/$(basename "$f")" "$f"; done
  if sshd -t 2>/dev/null; then
    reload_ssh || say "restore: reload failed"
    say "restore: sshd config put back from $BK"
  else
    say "restore: sshd -t failed even after restore; NOT reloading (running sshd keeps its old config)"
  fi
}

reload_ssh() {
  unit=ssh
  systemctl cat ssh.service >/dev/null 2>&1 || unit=sshd
  if systemctl is-active --quiet "$unit"; then
    systemctl reload "$unit"
  else
    say "$unit.service is not active (socket activation?); new connections read the new config"
    return 0
  fi
}

# 5. Edit.
for f in "$MAIN" $confs; do
  sed -i -E 's/^[[:space:]]*#?[[:space:]]*(PasswordAuthentication)[[:space:]]+.*/\1 no/I' "$f"
done
# Drop any earlier fillrate block, then make sure the main file sets it globally.
tmp="$STATE/sshd_config.new"
awk -v b="$BEGIN_MARK" -v e="$END_MARK" '$0 == b { skip = 1; next } $0 == e { skip = 0; next } !skip { print }' "$MAIN" > "$tmp"
if ! awk 'tolower($1) == "match" { exit 1 } tolower($1) == "passwordauthentication" { found = 1 } END { exit !found }' "$tmp"; then
  awk 'BEGIN { done = 0 } !done && tolower($1) == "match" { print "PasswordAuthentication no"; done = 1 } { print } END { if (!done) print "PasswordAuthentication no" }' "$tmp" > "$tmp.2" && mv "$tmp.2" "$tmp"
fi
{
  cat "$tmp"
  printf '%s\n' "$BEGIN_MARK"
  printf 'Match Address %s\n' "$ALLOW_CSV"
  printf '    PasswordAuthentication yes\n'
  printf '%s\n' "$END_MARK"
} > "$tmp.2"
cat "$tmp.2" > "$MAIN"
rm -f "$tmp" "$tmp.2"

# 6. Test, then reload.
if ! sshd -t; then
  restore_ssh
  say "summary: ssh-password-off result=sshd-t-failed restored=yes log=$LOG"
  exit 2
fi
reload_ssh || fail_verify "reload-failed"
say "sshd: config valid, reloaded (no restart)"

# 7. Verify.
after=$(sshd -T 2>/dev/null | awk '$1 == "passwordauthentication" { print $2; exit }')
say "after: passwordauthentication $after"
[ "$after" = "no" ] || fail_verify "global-still-$after"
for ip in $ALLOW; do
  v=$(sshd -T -C "addr=$ip,user=x,host=x" 2>/dev/null | awk '$1 == "passwordauthentication" { print $2; exit }')
  [ "$v" = "yes" ] || fail_verify "allowlisted-not-yes"
done
v=$(sshd -T -C "addr=$PROBE_OUTSIDE,user=x,host=x" 2>/dev/null | awk '$1 == "passwordauthentication" { print $2; exit }')
[ "$v" = "no" ] || fail_verify "outside-addr-not-no"
say "verify: global no, every allowlisted address yes, outside address no"

# 8. Firewall, add-only.
added=0; present=0
case "$fw" in
  ufw)
    for ip in $ALLOW; do for p in $ports; do
      out=$(ufw allow proto tcp from "$ip" to any port "$p" comment "$TAG" 2>&1)
      case "$out" in *existing*) present=$((present + 1)) ;; *) added=$((added + 1)) ;; esac
    done; done
    ;;
  iptables)
    for ip in $ALLOW; do for p in $ports; do
      if iptables -C INPUT -p tcp -s "$ip" --dport "$p" -m comment --comment "$TAG" -j ACCEPT 2>/dev/null; then
        present=$((present + 1))
      else
        iptables -I INPUT 1 -p tcp -s "$ip" --dport "$p" -m comment --comment "$TAG" -j ACCEPT && added=$((added + 1))
      fi
    done; done
    say "note: iptables rules are not saved; they last until reboot (nothing persisted, Docker rules untouched)"
    ;;
  nft)
    for ip in $ALLOW; do for p in $ports; do
      if nft list chain inet filter input 2>/dev/null | grep -F "ip saddr $ip tcp dport $p" | grep -q "$TAG"; then
        present=$((present + 1))
      else
        nft insert rule inet filter input ip saddr "$ip" tcp dport "$p" accept comment "\"$TAG\"" && added=$((added + 1))
      fi
    done; done
    ;;
  *)
    say "firewall: no ufw (active), iptables, or nft inet filter input chain; no rules added"
    ;;
esac
say "firewall: added=$added already_present=$present via=$fw"

# 9. Rollback.
say "rollback (run as root):"
say "  cp -p $BK/ssh/sshd_config $MAIN"
for f in $confs; do say "  cp -p $BK/ssh/sshd_config.d/$(basename "$f") $f"; done
say "  sshd -t && systemctl reload ssh"
case "$fw" in
  ufw) say "  ufw status numbered, then ufw delete <n> for each rule commented $TAG" ;;
  iptables) say "  for each: iptables -D INPUT -p tcp -s <ip> --dport <port> -m comment --comment $TAG -j ACCEPT" ;;
  nft) say "  nft -a list chain inet filter input, then nft delete rule inet filter input handle <h> for each $TAG rule" ;;
esac
say "  (the original pre-change files are also kept as <file>.orig)"

say "summary: ssh-password-off result=ok before=$before after=$after allowlist_count=$n firewall=$fw added=$added present=$present ufw=$ufw_state backup=$BK log=$LOG"
exit 0