#!/bin/sh
set -eu

HELPER_NAME="ShellOrchestra SSH key setup helper"
CA_FILE="/etc/ssh/shellorchestra_user_ca.pub"
CA_CONF_DIR="/etc/ssh/sshd_config.d"
CA_CONF_FILE="$CA_CONF_DIR/99-shellorchestra-user-ca.conf"
MAIN_SSHD_CONFIG="/etc/ssh/sshd_config"

info() {
  printf '%s: %s\n' "$HELPER_NAME" "$*" >&2
}

fail() {
  printf '%s error: %s\n' "$HELPER_NAME" "$*" >&2
  exit 1
}

usage() {
  cat >&2 <<'USAGE'
ShellOrchestra SSH key setup helper

Usage:
  sh install-posix.sh <CA_KEY_B64URL>
  sh install-posix.sh --ca <CA_KEY_B64URL>
  sh install-posix.sh --classic [--account USER] <AUTHORIZED_KEYS_LINE_B64URL>

The helper detects the supported POSIX platform before changing SSH
configuration. The default mode writes a ShellOrchestra SSH CA public key and
configures OpenSSH TrustedUserCAKeys. Classic mode writes one permanent
authorized_keys line for one account and should be used only when the server
deliberately does not use SSH CA certificates.
USAGE
}

run_root() {
  if [ "$(id -u)" -eq 0 ]; then
    "$@"
    return
  fi
  if ! command -v sudo >/dev/null 2>&1; then
    fail "root privileges are required, and sudo is not installed"
  fi
  sudo "$@"
}

mode="ca"
target_account="${SHELLORCHESTRA_TARGET_USER:-root}"
encoded_payload=""

while [ "$#" -gt 0 ]; do
  case "$1" in
    --ca)
      mode="ca"
      shift
      ;;
    --classic)
      mode="classic"
      shift
      ;;
    --account)
      shift
      [ "$#" -gt 0 ] || fail "--account requires a user name"
      target_account="$1"
      shift
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    --*)
      fail "unsupported option: $1"
      ;;
    *)
      [ -z "$encoded_payload" ] || fail "only one encoded key payload is allowed"
      encoded_payload="$1"
      shift
      ;;
  esac
done

[ -n "$encoded_payload" ] || fail "encoded key payload is required"
[ -n "$target_account" ] || fail "target account cannot be empty"

base64_decode() {
  if command -v base64 >/dev/null 2>&1; then
    case "$(uname -s)" in
      Darwin|FreeBSD|OpenBSD|NetBSD)
        base64 -D
        ;;
      *)
        base64 -d
        ;;
    esac
    return
  fi
  fail "base64 command is required to decode the key payload"
}

decode_base64url() {
  value="$1"
  case "$value" in
    *[!A-Za-z0-9_-]*)
      fail "encoded key payload contains characters outside base64url"
      ;;
  esac
  length_mod=$((${#value} % 4))
  case "$length_mod" in
    0) padding="" ;;
    2) padding="==" ;;
    3) padding="=" ;;
    *) fail "encoded key payload has invalid base64url length" ;;
  esac
  printf '%s%s' "$value" "$padding" | tr '_-' '/+' | base64_decode
}

payload=$(decode_base64url "$encoded_payload" | tr -d '\r')
payload=$(printf '%s' "$payload" | sed 's/[[:space:]]*$//')
[ -n "$payload" ] || fail "decoded key payload is empty"

validate_ca_key() {
  set -- $1
  [ "$#" -ge 2 ] || fail "decoded CA key is not an OpenSSH public key"
  [ "$1" = "ssh-ed25519" ] || fail "ShellOrchestra currently expects an ssh-ed25519 public key"
}

validate_authorized_key_line() {
  line="$1"
  case "$line" in
    ssh-ed25519\ *|from=\"*\"\ ssh-ed25519\ *)
      return
      ;;
  esac
  fail "decoded classic payload must be an ssh-ed25519 authorized_keys line"
}

linux_sshd_reload() {
  if command -v systemctl >/dev/null 2>&1; then
    if systemctl list-unit-files sshd.service >/dev/null 2>&1; then
      run_root systemctl reload sshd
      return
    fi
    if systemctl list-unit-files ssh.service >/dev/null 2>&1; then
      run_root systemctl reload ssh
      return
    fi
  fi
  if command -v service >/dev/null 2>&1; then
    if service sshd status >/dev/null 2>&1; then
      run_root service sshd reload
      return
    fi
    if service ssh status >/dev/null 2>&1; then
      run_root service ssh reload
      return
    fi
  fi
  fail "OpenSSH configuration is valid, but SSH reload method was not detected. Reload ssh or sshd manually."
}

install_ca_linux() {
  tmp_file=$(mktemp)
  printf '%s\n' "$payload" > "$tmp_file"
  run_root install -d -o root -g root -m 0755 "$(dirname "$CA_FILE")"
  run_root install -o root -g root -m 0644 "$tmp_file" "$CA_FILE"
  rm -f "$tmp_file"

  if [ -d "$CA_CONF_DIR" ] || grep -Eq '^[[:space:]]*Include[[:space:]]+/etc/ssh/sshd_config\.d/\*\.conf' "$MAIN_SSHD_CONFIG" 2>/dev/null; then
    run_root install -d -o root -g root -m 0755 "$CA_CONF_DIR"
    printf 'TrustedUserCAKeys %s\n' "$CA_FILE" | run_root tee "$CA_CONF_FILE" >/dev/null
  elif grep -Eq '^[[:space:]]*TrustedUserCAKeys[[:space:]]+' "$MAIN_SSHD_CONFIG" 2>/dev/null; then
    fail "TrustedUserCAKeys is already configured in $MAIN_SSHD_CONFIG. Update it manually or remove the old directive before using this helper."
  else
    printf '\n# ShellOrchestra SSH CA\nTrustedUserCAKeys %s\n' "$CA_FILE" | run_root tee -a "$MAIN_SSHD_CONFIG" >/dev/null
  fi

  run_root sshd -t
  linux_sshd_reload
}

install_ca_macos() {
  tmp_file=$(mktemp)
  printf '%s\n' "$payload" > "$tmp_file"
  run_root install -o root -g wheel -m 0644 "$tmp_file" "$CA_FILE"
  rm -f "$tmp_file"

  if ! grep -Eq '^[[:space:]]*TrustedUserCAKeys[[:space:]]+/etc/ssh/shellorchestra_user_ca\.pub[[:space:]]*$' "$MAIN_SSHD_CONFIG" 2>/dev/null; then
    if grep -Eq '^[[:space:]]*TrustedUserCAKeys[[:space:]]+' "$MAIN_SSHD_CONFIG" 2>/dev/null; then
      fail "TrustedUserCAKeys is already configured in $MAIN_SSHD_CONFIG. Update it manually or remove the old directive before using this helper."
    fi
    printf '\n# ShellOrchestra SSH CA\nTrustedUserCAKeys %s\n' "$CA_FILE" | run_root tee -a "$MAIN_SSHD_CONFIG" >/dev/null
  fi

  run_root /usr/sbin/sshd -t -f "$MAIN_SSHD_CONFIG"
  run_root launchctl kickstart -k system/com.openssh.sshd
}

account_home_linux() {
  if [ "$target_account" = "root" ]; then
    printf '/root\n'
    return
  fi
  if command -v getent >/dev/null 2>&1; then
    home=$(getent passwd "$target_account" | awk -F: '{print $6}')
    [ -n "$home" ] || fail "target account was not found: $target_account"
    printf '%s\n' "$home"
    return
  fi
  fail "getent is required to locate non-root target accounts on this platform"
}

account_home_macos() {
  if [ "$target_account" = "root" ]; then
    printf '/var/root\n'
    return
  fi
  home=$(dscl . -read "/Users/$target_account" NFSHomeDirectory 2>/dev/null | awk '{print $2}')
  [ -n "$home" ] || fail "target account was not found: $target_account"
  printf '%s\n' "$home"
}

install_classic_authorized_key() {
  home_dir="$1"
  key_line="$payload"
  group_name=$(id -gn "$target_account")
  ssh_dir="$home_dir/.ssh"
  authorized_keys="$ssh_dir/authorized_keys"

  run_root install -d -o "$target_account" -g "$group_name" -m 0700 "$ssh_dir"
  run_root touch "$authorized_keys"
  run_root chown "$target_account:$group_name" "$authorized_keys"
  run_root chmod 0600 "$authorized_keys"
  if ! run_root grep -qxF "$key_line" "$authorized_keys" 2>/dev/null; then
    printf '%s\n' "$key_line" | run_root tee -a "$authorized_keys" >/dev/null
  fi
}

case "$mode" in
  ca)
    validate_ca_key "$payload"
    case "$(uname -s)" in
      Linux)
        info "Installing ShellOrchestra SSH CA public key for Linux OpenSSH"
        install_ca_linux
        ;;
      Darwin)
        info "Installing ShellOrchestra SSH CA public key for macOS OpenSSH"
        install_ca_macos
        ;;
      *)
        fail "unsupported operating system for CA installation: $(uname -s)"
        ;;
    esac
    info "Done. This server now trusts ShellOrchestra short-lived SSH certificates."
    ;;
  classic)
    validate_authorized_key_line "$payload"
    case "$(uname -s)" in
      Linux)
        info "Installing ShellOrchestra classic fallback key for account: $target_account"
        install_classic_authorized_key "$(account_home_linux)"
        ;;
      Darwin)
        info "Installing ShellOrchestra classic fallback key for account: $target_account"
        install_classic_authorized_key "$(account_home_macos)"
        ;;
      *)
        fail "unsupported operating system for classic fallback installation: $(uname -s)"
        ;;
    esac
    info "Done. Classic permanent-key fallback was installed for $target_account."
    ;;
  *)
    fail "unsupported install mode: $mode"
    ;;
esac