#!/usr/bin/env bash
set -Eeuo pipefail

# Tests may prepend fake binaries. Production always uses a deterministic PATH.
if [[ "${TEST_MODE:-0}" == "1" ]]; then
  :
else
  export PATH=/usr/sbin:/usr/bin:/sbin:/bin
fi

AI_USER="${AI_USER:-aiops}"
AI_ALLOWED_FROM="${AI_ALLOWED_FROM:-11.57.0.20/32}"
COMPOSE_PROJECT="${COMPOSE_PROJECT:-qnap}"
PUBLIC_KEY_FILE="${PUBLIC_KEY_FILE:-}"
CONFIG_DIR="/etc/ai-readonly"
RUNTIME_COMMAND="/usr/local/libexec/ai-readonly-command"
SSHD_DROPIN="/etc/ssh/sshd_config.d/91-ai-readonly.conf"
HARDENING_DROPIN="/etc/ssh/sshd_config.d/99-hardening.conf"
SUDOERS_FILE="/etc/sudoers.d/ai-readonly"
MANAGED_MARKER="$CONFIG_DIR/managed"

if [[ "${TEST_MODE:-0}" == "1" ]]; then
  if [[ -n "${TEST_CONFIG_DIR:-}" ]]; then
    CONFIG_DIR="$TEST_CONFIG_DIR"
    MANAGED_MARKER="$CONFIG_DIR/managed"
  fi
  if [[ -n "${TEST_RUNTIME_DIR:-}" ]]; then
    RUNTIME_COMMAND="$TEST_RUNTIME_DIR/ai-readonly-command"
  fi
  if [[ -n "${TEST_SSHD_DROPIN:-}" ]]; then
    SSHD_DROPIN="$TEST_SSHD_DROPIN"
  fi
  if [[ -n "${TEST_SUDOERS_FILE:-}" ]]; then
    SUDOERS_FILE="$TEST_SUDOERS_FILE"
  fi
fi

SCRIPT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SOURCE_COMMAND="$SCRIPT_ROOT/ai-readonly-command.sh"

UNINSTALL=0
ALLOW_ANY_SOURCE=0
ACK_ANY_SOURCE=0
DISCOVER_CONTAINERS=0
ROTATE_PUBLIC_KEY=0
EXTRA_CONTAINERS=()
EXTRA_UNITS=()

INSTALLED_RUNTIME=0
INSTALLED_ALLOWLISTS=0
INSTALLED_SSHD_DROPIN=0
INSTALLED_SUDOERS=0
INSTALLED_AUTHORIZED_KEYS=0
INSTALLED_MARKER=0
INSTALL_COMMITTED=0
INSTALL_ROLLED_BACK=0
existing_install=0
original_created_user=0
created_user=0
user_home=""
SNAPSHOT_DIR=""
SNAPSHOT_COMPLETE=0
SSHD_RELOAD_ATTEMPTED=0

log() { printf '\n\033[1;34m==> %s\033[0m\n' "$*"; }
warn() { printf '\n\033[1;33mWARNING: %s\033[0m\n' "$*" >&2; }
die() {
  printf '\n\033[1;31mERROR: %s\033[0m\n' "$*" >&2
  exit 1
}

expected_home_for() {
  if [[ "${TEST_MODE:-0}" == "1" ]]; then
    printf '%s/home/%s\n' "${TEST_TMP:-/tmp}" "$1"
  else
    printf '/home/%s\n' "$1"
  fi
}

marker_field() {
  [[ -f "$2" ]] || return 0
  awk -F= -v key="$1" '$1 == key { print $2; exit }' "$2"
}

reload_sshd() {
  if systemctl is-active --quiet ssh.service 2>/dev/null; then
    SSHD_RELOAD_ATTEMPTED=1
    systemctl reload ssh.service
  elif systemctl is-active --quiet sshd.service 2>/dev/null; then
    SSHD_RELOAD_ATTEMPTED=1
    systemctl reload sshd.service
  else
    warn "SSH service is not active; configuration was installed but not reloaded."
  fi
}

SNAPSHOT_NAMES=(runtime containers units dropin sudoers authorized_keys marker)

snapshot_path_for() {
  case "$1" in
    runtime) printf '%s\n' "$RUNTIME_COMMAND" ;;
    containers) printf '%s\n' "$CONFIG_DIR/containers.allow" ;;
    units) printf '%s\n' "$CONFIG_DIR/units.allow" ;;
    dropin) printf '%s\n' "$SSHD_DROPIN" ;;
    sudoers) printf '%s\n' "$SUDOERS_FILE" ;;
    authorized_keys)
      [[ -n "$user_home" ]] || return 0
      printf '%s\n' "$user_home/.ssh/authorized_keys"
      ;;
    marker) printf '%s\n' "$MANAGED_MARKER" ;;
  esac
}

# Publish SNAPSHOT_DIR only after every copy succeeded and a completion marker
# exists. A failed snapshot therefore cannot be mistaken for a complete image.
create_snapshot() {
  local tmp name path
  tmp="$(mktemp -d)" || return 1
  chmod 0700 "$tmp" || {
    rm -rf -- "$tmp" 2>/dev/null || true
    return 1
  }

  for name in "${SNAPSHOT_NAMES[@]}"; do
    path="$(snapshot_path_for "$name")"
    [[ -n "$path" && -e "$path" ]] || continue
    if ! cp -a -- "$path" "$tmp/$name"; then
      rm -rf -- "$tmp" 2>/dev/null || true
      return 1
    fi
  done

  if ! : >"$tmp/.complete"; then
    rm -rf -- "$tmp" 2>/dev/null || true
    return 1
  fi
  chmod 0600 "$tmp/.complete" || {
    rm -rf -- "$tmp" 2>/dev/null || true
    return 1
  }

  SNAPSHOT_DIR="$tmp"
  SNAPSHOT_COMPLETE=1
}

restore_snapshot() {
  [[ "$SNAPSHOT_COMPLETE" == "1" ]] || return 1
  [[ -n "$SNAPSHOT_DIR" && -d "$SNAPSHOT_DIR" && -f "$SNAPSHOT_DIR/.complete" ]] || return 1

  local rc=0 name path restore_tmp
  for name in "${SNAPSHOT_NAMES[@]}"; do
    path="$(snapshot_path_for "$name")"
    [[ -n "$path" ]] || continue

    if [[ -e "$SNAPSHOT_DIR/$name" ]]; then
      mkdir -p "$(dirname "$path")" 2>/dev/null || {
        rc=1
        continue
      }
      restore_tmp="${path}.restore.$$"
      rm -f -- "$restore_tmp" 2>/dev/null || true
      if cp -a -- "$SNAPSHOT_DIR/$name" "$restore_tmp" && mv -f -- "$restore_tmp" "$path"; then
        :
      else
        rm -f -- "$restore_tmp" 2>/dev/null || true
        rc=1
      fi
    else
      rm -f -- "$path" 2>/dev/null || rc=1
    fi
  done
  return "$rc"
}

rollback_install() {
  local rc=$?
  [[ "$INSTALL_COMMITTED" == "1" ]] && return 0
  if [[ "$INSTALL_ROLLED_BACK" == "1" ]]; then
    trap - EXIT
    exit "$rc"
  fi
  INSTALL_ROLLED_BACK=1

  local restore_failed=0
  if [[ "$existing_install" == "1" ]]; then
    if [[ "$SNAPSHOT_COMPLETE" == "1" && -f "${SNAPSHOT_DIR:-}/.complete" ]]; then
      warn "Restoring the previous managed installation (rc=$rc)"
      restore_snapshot || restore_failed=1
    else
      # Snapshot creation is the last operation before any managed mutation.
      # If it did not complete, the original installation is still untouched.
      warn "Pre-update snapshot did not complete; the existing managed installation was left untouched (rc=$rc)"
    fi
  else
    warn "Rolling back partial install (rc=$rc)"
    if [[ "$INSTALLED_MARKER" == "1" ]]; then
      rm -f "$MANAGED_MARKER" 2>/dev/null || true
    fi
    if [[ "$INSTALLED_AUTHORIZED_KEYS" == "1" && -n "$user_home" ]]; then
      rm -f "$user_home/.ssh/authorized_keys" 2>/dev/null || true
    fi
    if [[ "$INSTALLED_SUDOERS" == "1" ]]; then
      rm -f "$SUDOERS_FILE" 2>/dev/null || true
    fi
    if [[ "$INSTALLED_SSHD_DROPIN" == "1" ]]; then
      rm -f "$SSHD_DROPIN" 2>/dev/null || true
    fi
    if [[ "$INSTALLED_RUNTIME" == "1" ]]; then
      rm -f "$RUNTIME_COMMAND" 2>/dev/null || true
    fi
    if [[ "$INSTALLED_ALLOWLISTS" == "1" ]]; then
      rm -rf "$CONFIG_DIR" 2>/dev/null || true
    fi

    if [[ "$created_user" == "1" ]] && id "$AI_USER" >/dev/null 2>&1; then
      userdel --remove "$AI_USER" >/dev/null 2>&1 \
        || userdel "$AI_USER" >/dev/null 2>&1 \
        || true
    fi
  fi

  if [[ "$SSHD_RELOAD_ATTEMPTED" == "1" ]]; then
    if "$SSHD_BIN" -t; then
      reload_sshd || warn "SSH reload failed during rollback; inspect the SSH service before retrying."
    else
      warn "sshd -t failed after rollback; keep the administrator session open and inspect /etc/ssh/sshd_config*."
      restore_failed=1
    fi
  fi

  if ((restore_failed == 1)); then
    warn "Rollback could not fully reconcile the host. The snapshot is retained at ${SNAPSHOT_DIR:-<none>} for manual recovery."
  elif [[ -n "$SNAPSHOT_DIR" ]]; then
    rm -rf -- "$SNAPSHOT_DIR" 2>/dev/null || true
  fi

  trap - EXIT
  exit "$rc"
}

usage() {
  cat <<'USAGE'
Install a forced-command, read-only SSH account for AI/automation diagnostics.

Usage:
  sudo ./scripts/install-ai-readonly-ssh.sh --public-key-file PATH [options]
  sudo ./scripts/install-ai-readonly-ssh.sh --uninstall [--user NAME]

Options:
  --public-key-file PATH   File containing exactly one SSH public key.
  --rotate-public-key      Explicitly replace the key of an existing managed
                           installation. Without this flag, a different key is
                           rejected. Not valid for a fresh installation.
  --allowed-from CIDRS     Comma-separated IPv4/IPv6 addresses or CIDRs.
                           Default: 11.57.0.20/32 (one WireGuard peer).
  --allow-any-source       Omit the authorized_keys source restriction.
                           Requires --acknowledge-any-source. Risky.
  --acknowledge-any-source Acknowledge the weakened source restriction.
  --user NAME              Dedicated account name. Default: aiops.
  --compose-project NAME   Compose project for optional discovery. Default: hrvlo.
  --discover-containers    Opt in to docker ps container discovery. Off by default.
  --container NAME         Add an exact container name. Repeatable.
  --unit UNIT              Add an exact systemd unit. Repeatable.
  --uninstall              Remove files managed by this installer.
  -h, --help               Show this help.

Examples:
  sudo ./scripts/install-ai-readonly-ssh.sh \
    --public-key-file /root/ai-reader.pub \
    --allowed-from 11.57.0.20/32

  sudo ./scripts/install-ai-readonly-ssh.sh \
    --public-key-file /root/ai-reader-next.pub \
    --rotate-public-key \
    --allowed-from 11.57.0.20/32
USAGE
}

while (($# > 0)); do
  case "$1" in
    --public-key-file)
      (($# >= 2)) || die "--public-key-file requires a path"
      PUBLIC_KEY_FILE="$2"
      shift 2
      ;;
    --rotate-public-key)
      ROTATE_PUBLIC_KEY=1
      shift
      ;;
    --allowed-from)
      (($# >= 2)) || die "--allowed-from requires a value"
      AI_ALLOWED_FROM="$2"
      shift 2
      ;;
    --allow-any-source)
      ALLOW_ANY_SOURCE=1
      shift
      ;;
    --acknowledge-any-source)
      ACK_ANY_SOURCE=1
      shift
      ;;
    --user)
      (($# >= 2)) || die "--user requires a value"
      AI_USER="$2"
      shift 2
      ;;
    --compose-project)
      (($# >= 2)) || die "--compose-project requires a value"
      COMPOSE_PROJECT="$2"
      shift 2
      ;;
    --discover-containers)
      DISCOVER_CONTAINERS=1
      shift
      ;;
    --container)
      (($# >= 2)) || die "--container requires a value"
      EXTRA_CONTAINERS+=("$2")
      shift 2
      ;;
    --unit)
      (($# >= 2)) || die "--unit requires a value"
      EXTRA_UNITS+=("$2")
      shift 2
      ;;
    --uninstall)
      UNINSTALL=1
      shift
      ;;
    -h | --help)
      usage
      exit 0
      ;;
    *) die "unknown argument: $1" ;;
  esac
done

TEST_MODE="${TEST_MODE:-0}"
if [[ "$TEST_MODE" != "1" ]]; then
  ((EUID == 0)) || die "run this installer with sudo or as root"
  [[ "$(uname -s)" == "Linux" ]] || die "Linux is required"
fi
[[ "$AI_USER" =~ ^[a-z_][a-z0-9_-]{0,30}$ ]] || die "invalid account name: $AI_USER"
[[ "$COMPOSE_PROJECT" =~ ^[A-Za-z0-9][A-Za-z0-9_.-]{0,63}$ ]] || die "invalid Compose project name"

if ((ALLOW_ANY_SOURCE == 1)) && ((ACK_ANY_SOURCE == 0)); then
  die "--allow-any-source requires --acknowledge-any-source"
fi
if ((UNINSTALL == 1)) && ((ROTATE_PUBLIC_KEY == 1)); then
  die "--rotate-public-key cannot be combined with --uninstall"
fi

SSHD_BIN="$(command -v sshd 2>/dev/null || true)"
[[ -n "$SSHD_BIN" ]] || SSHD_BIN=/usr/sbin/sshd

restore_staged_dropin() {
  local staged="$1"
  [[ -n "$staged" && -e "$staged" ]] || return 0
  install -d -o root -g root -m 0755 "$(dirname "$SSHD_DROPIN")" 2>/dev/null || true
  rm -f -- "$SSHD_DROPIN" 2>/dev/null || true
  cp -a -- "$staged" "$SSHD_DROPIN"
}

uninstall_access() {
  [[ -f "$MANAGED_MARKER" ]] || die "managed marker not found; refusing to remove an unknown installation"
  # shellcheck disable=SC1090
  source "$MANAGED_MARKER"
  [[ "${MANAGED_AI_USER:-}" == "$AI_USER" ]] \
    || die "managed account is ${MANAGED_AI_USER:-unknown}; rerun with --user ${MANAGED_AI_USER:-<name>}"

  log "Validating SSH configuration without the managed drop-in"
  local stage="" staged_dropin=""
  if [[ -e "$SSHD_DROPIN" ]]; then
    stage="$(mktemp -d)"
    chmod 0700 "$stage"
    staged_dropin="$stage/dropin"
    cp -a -- "$SSHD_DROPIN" "$staged_dropin" \
      || die "could not stage the SSH drop-in; no managed artifact was removed"
    rm -f -- "$SSHD_DROPIN"
  fi

  install -d -o root -g root -m 0755 /run/sshd
  if ! "$SSHD_BIN" -t; then
    if [[ -n "$staged_dropin" ]]; then
      if restore_staged_dropin "$staged_dropin" && "$SSHD_BIN" -t; then
        rm -rf -- "$stage"
      else
        warn "Could not restore ${SSHD_DROPIN}; the staged copy is retained at ${staged_dropin}."
      fi
    fi
    die "sshd rejected the configuration without ${SSHD_DROPIN}; no managed artifact was removed"
  fi

  if ! reload_sshd; then
    warn "SSH reload failed; restoring ${SSHD_DROPIN} so the running daemon and filesystem agree."
    if restore_staged_dropin "$staged_dropin" && "$SSHD_BIN" -t; then
      [[ -n "$stage" ]] && rm -rf -- "$stage"
    else
      warn "Could not restore ${SSHD_DROPIN}; the staged copy is retained at ${staged_dropin:-<none>}."
    fi
    reload_sshd || warn "SSH reload failed again; inspect the SSH service before retrying --uninstall."
    die "uninstall aborted before removing managed artifacts; the managed marker was retained"
  fi

  [[ -n "$stage" ]] && rm -rf -- "$stage"
  log "Removing managed sudo, runtime and account artifacts"
  rm -f -- "$SUDOERS_FILE" "$RUNTIME_COMMAND"

  user_home="$(getent passwd "$AI_USER" | cut -d: -f6 || true)"
  if [[ -n "$user_home" && "$user_home" == "$(expected_home_for "$AI_USER")" && "${CREATED_USER:-0}" == "1" ]]; then
    rm -rf -- "$user_home"
  fi
  if [[ "${CREATED_USER:-0}" == "1" ]] && id "$AI_USER" >/dev/null 2>&1; then
    userdel --remove "$AI_USER" >/dev/null 2>&1 || userdel "$AI_USER" >/dev/null 2>&1 || true
  fi

  rm -rf -- "$CONFIG_DIR"
  log "Read-only AI SSH access removed"
}

if ((UNINSTALL == 1)); then
  uninstall_access
  exit 0
fi

[[ -n "$PUBLIC_KEY_FILE" ]] || die "--public-key-file is required"
[[ -r "$PUBLIC_KEY_FILE" ]] || die "cannot read public key file: $PUBLIC_KEY_FILE"
[[ -f "$SOURCE_COMMAND" ]] || die "missing companion command: $SOURCE_COMMAND"

for cmd in install useradd userdel usermod passwd sudo visudo stat grep sort awk sed systemctl ssh-keygen getent gpasswd mktemp tr cut flock find python3 cp mv rm mkdir chmod chown id docker; do
  command -v "$cmd" >/dev/null || die "required command is missing: $cmd"
done
[[ -x "$SSHD_BIN" ]] || die "sshd is not installed"

nonempty_key_lines="$(grep -cvE '^[[:space:]]*(#|$)' "$PUBLIC_KEY_FILE" || true)"
[[ "$nonempty_key_lines" == "1" ]] || die "public key file must contain exactly one non-comment key"
public_key="$(grep -vE '^[[:space:]]*(#|$)' "$PUBLIC_KEY_FILE")"
read -r key_type key_blob _ <<<"$public_key"
case "$key_type" in
  ssh-ed25519 | ssh-rsa | ecdsa-sha2-nistp256 | ecdsa-sha2-nistp384 | ecdsa-sha2-nistp521 | [email protected] | [email protected]) ;;
  *) die "unsupported or malformed SSH public key type: $key_type" ;;
esac
[[ "$key_blob" =~ ^[A-Za-z0-9+/]+={0,3}$ ]] || die "malformed SSH public key data"

key_tmp="$(mktemp)"
trap 'rollback_install; rm -f "$key_tmp"' EXIT
printf '%s\n' "$public_key" >"$key_tmp"
ssh-keygen -l -f "$key_tmp" >/dev/null 2>&1 || die "ssh-keygen could not validate the public key"

if ((ALLOW_ANY_SOURCE == 0)); then
  [[ -n "$AI_ALLOWED_FROM" ]] || die "--allowed-from must not be empty"
  [[ "$AI_ALLOWED_FROM" =~ ^[0-9A-Fa-f:.,/]+$ ]] \
    || die "--allowed-from accepts only comma-separated IP addresses or CIDRs"
  python3 - "$AI_ALLOWED_FROM" <<'PY'
import ipaddress
import sys
for value in sys.argv[1].split(','):
    if not value:
        raise SystemExit('empty source entry')
    ipaddress.ip_network(value, strict=False)
PY
fi

validate_container_name() {
  [[ "$1" =~ ^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$ ]] || die "invalid container name: $1"
}
validate_unit_name() {
  [[ "$1" =~ ^[A-Za-z0-9_.@-]+\.(service|socket|timer)$ ]] || die "invalid systemd unit: $1"
}
for item in "${EXTRA_CONTAINERS[@]}"; do validate_container_name "$item"; done
for item in "${EXTRA_UNITS[@]}"; do validate_unit_name "$item"; done

if [[ -f "$MANAGED_MARKER" ]]; then
  existing_install=1
  original_created_user="$(marker_field CREATED_USER "$MANAGED_MARKER")"
  [[ "$original_created_user" == "1" ]] || original_created_user=0
fi
if ((ROTATE_PUBLIC_KEY == 1)) && ((existing_install == 0)); then
  die "--rotate-public-key requires an existing managed installation"
fi

require_safe_target() {
  local target="$1" managed_ai_user="" marker_created="" target_uid existing_keys key_mode line_count
  if [[ -f "$MANAGED_MARKER" ]]; then
    managed_ai_user="$(marker_field MANAGED_AI_USER "$MANAGED_MARKER")"
    marker_created="$(marker_field CREATED_USER "$MANAGED_MARKER")"
  fi

  if [[ -n "$managed_ai_user" && "$managed_ai_user" != "$target" ]]; then
    die "managed marker records MANAGED_AI_USER=${managed_ai_user}; rerun with --user ${managed_ai_user} or --uninstall first"
  fi

  if id "$target" >/dev/null 2>&1; then
    if [[ -z "$managed_ai_user" || "$marker_created" != "1" ]]; then
      die "account $target already exists but is not marked as created by this installer; choose another --user"
    fi
  else
    return 0
  fi

  target_uid="$(id -u "$target")"
  ((target_uid >= 1000)) || die "target account $target has UID $target_uid (< 1000); refusing to modify a system or human user"
  for group in sudo admin adm docker systemd-journal; do
    if id -nG "$target" | tr ' ' '\n' | grep -Fxq "$group"; then
      die "target account $target is in group $group; refusing to take over a privileged account"
    fi
  done
  if [[ -n "${SUDO_USER:-}" && "$target" == "$SUDO_USER" ]]; then
    die "refusing to modify the invoking administrator (${SUDO_USER})"
  fi

  existing_keys="$(expected_home_for "$target")/.ssh/authorized_keys"
  if [[ -e "$existing_keys" ]]; then
    [[ -f "$existing_keys" && ! -L "$existing_keys" ]] || die "managed authorized_keys must be a regular file"
    key_mode="$(stat -c '%U:%G:%a' "$existing_keys" 2>/dev/null || true)"
    case "$key_mode" in
      root:root:644 | root:root:640 | root:root:600) ;;
      *) die "unsafe managed authorized_keys ownership or mode: ${key_mode:-unknown}" ;;
    esac
    line_count="$(grep -cvE '^[[:space:]]*(#|$)' "$existing_keys" || true)"
    [[ "$line_count" == "1" ]] || die "managed authorized_keys must contain exactly one key"
    grep -Fq "restrict,command=\"${RUNTIME_COMMAND}\"" "$existing_keys" \
      || die "authorized_keys is not managed by this installer"

    if ! grep -Fq -- "$key_blob" "$existing_keys"; then
      if ((ROTATE_PUBLIC_KEY == 1)); then
        warn "Rotating the public key for managed account $target"
      else
        die "supplied public key differs from the managed key; rerun with --rotate-public-key to rotate it explicitly"
      fi
    fi
  fi
}
require_safe_target "$AI_USER"

fail_on_allowlist() {
  local effective="$1" directive value admits token
  for directive in allowusers denyusers allowgroups denygroups; do
    # OpenSSH 9.4+ emits each entry of a multi-value directive on its own
    # line (e.g. "allowusers alice" / "allowusers bob"). Older releases
    # joined them with spaces. Gather every matching line and collapse them
    # into a single space-separated list so the token check below stays
    # correct on both formats.
    value="$(awk -v key="$directive" 'tolower($1) == key {$1=""; sub(/^ /, ""); print}' <<<"$effective" | paste -sd' ' -)"
    case "$directive" in
      allowusers)
        if [[ -n "$value" ]]; then
          admits=0
          for token in $value; do
            if [[ "$token" == "$AI_USER" || "$token" == "*" ]]; then
              admits=1
              break
            fi
          done
          ((admits == 1)) || die "effective AllowUsers is '${value}'. The account $AI_USER is not in the allowlist; update ${HARDENING_DROPIN} and run sudo sshd -t"
        fi
        ;;
      allowgroups)
        [[ -z "$value" ]] || warn "Effective AllowGroups is configured: ${value}; confirm $AI_USER is admitted."
        ;;
      denyusers | denygroups)
        [[ -z "$value" ]] || warn "Effective ${directive} is configured: ${value}; confirm $AI_USER is not denied."
        ;;
    esac
  done
}

log "Validating SSH allowlist for $AI_USER"
fail_on_allowlist "$("$SSHD_BIN" -T -C "user=${AI_USER},host=localhost,addr=11.57.0.20")"

if id "$AI_USER" >/dev/null 2>&1; then
  log "Updating existing managed account: $AI_USER"
else
  log "Creating dedicated account: $AI_USER"
  useradd --create-home --shell /bin/sh --user-group "$AI_USER"
  created_user=1
fi

user_home="$(getent passwd "$AI_USER" | cut -d: -f6)"
expected_home="$(expected_home_for "$AI_USER")"
[[ "$user_home" == "$expected_home" ]] || die "managed account home is $user_home; expected $expected_home"

if ((existing_install == 1)); then
  create_snapshot || die "could not create a complete pre-update snapshot; no managed file was changed"
fi

usermod --shell /bin/sh "$AI_USER"
passwd --lock "$AI_USER" >/dev/null
primary_group="$(id -gn "$AI_USER")"
while IFS= read -r supplementary_group; do
  [[ -n "$supplementary_group" && "$supplementary_group" != "$primary_group" ]] || continue
  gpasswd -d "$AI_USER" "$supplementary_group" >/dev/null
done < <(id -nG "$AI_USER" | tr ' ' '\n')

log "Installing root-owned command and allowlists"
install -d -o root -g root -m 0755 "$(dirname "$RUNTIME_COMMAND")" "$CONFIG_DIR"
install -o root -g root -m 0755 "$SOURCE_COMMAND" "$RUNTIME_COMMAND"
INSTALLED_RUNTIME=1
bash -n "$RUNTIME_COMMAND" || die "runtime command failed syntax check: $RUNTIME_COMMAND"

containers_tmp="$(mktemp)"
units_tmp="$(mktemp)"
sudoers_tmp="$(mktemp)"
sshd_tmp="$(mktemp)"
trap 'rollback_install; rm -f "$key_tmp" "$containers_tmp" "$units_tmp" "$sudoers_tmp" "$sshd_tmp"' EXIT

cat >"$containers_tmp" <<'CONTAINERS'
alloy
caddy
content-import-worker
db-backup
db-backup-r2
go-api
redis
vpn-dns
CONTAINERS

if ((DISCOVER_CONTAINERS == 1)); then
  warn "Auto-discovering containers; review ${CONFIG_DIR}/containers.allow after installation."
  while IFS= read -r name; do
    [[ -n "$name" ]] || continue
    case "$name" in
      supabase-auth | supabase-db | supabase-rest | supabase-realtime | supabase-storage | supabase-meta | supabase-functions | supabase-studio | supabase-kong | supabase-supavisor | supabase-imgproxy | supabase-edge-runtime | supabase-logflare | supabase-vector | supabase-analytics) ;;
      *)
        validate_container_name "$name"
        printf '%s\n' "$name" >>"$containers_tmp"
        ;;
    esac
  done < <(docker ps -a --filter "label=com.docker.compose.project=${COMPOSE_PROJECT}" --format '{{.Names}}' 2>/dev/null || true)
fi
printf '%s\n' "${EXTRA_CONTAINERS[@]}" >>"$containers_tmp"
sed -i '/^$/d' "$containers_tmp"
sort -u "$containers_tmp" -o "$containers_tmp"
install -o root -g root -m 0644 "$containers_tmp" "$CONFIG_DIR/containers.allow"
INSTALLED_ALLOWLISTS=1

cat >"$units_tmp" <<'UNITS'
cockpit.socket
docker.service
fail2ban.service
backup-checksums.service
backup-checksums.timer
bootstrap-ssh-cutoff.timer
docker-firewall.service
host-firewall.service
offhost-backup.service
postgres-globals.service
postgres-globals.timer
vnstat-dashboard.timer
UNITS
if systemctl list-unit-files ssh.service --no-legend 2>/dev/null | grep -q '^ssh.service'; then
  printf '%s\n' ssh.service >>"$units_tmp"
elif systemctl list-unit-files sshd.service --no-legend 2>/dev/null | grep -q '^sshd.service'; then
  printf '%s\n' sshd.service >>"$units_tmp"
fi
printf '%s\n' "${EXTRA_UNITS[@]}" >>"$units_tmp"
sed -i '/^$/d' "$units_tmp"
sort -u "$units_tmp" -o "$units_tmp"
install -o root -g root -m 0644 "$units_tmp" "$CONFIG_DIR/units.allow"

# Preserve SSH_CONNECTION, SSH_CLIENT, and SSH_ORIGINAL_COMMAND across sudo
# so the dispatcher's audit trail can attribute calls to the WireGuard peer
# (11.57.0.20) instead of falling back to source=local when sudo strips the
# env vars. Targeted via Defaults>aiops so the env_keep applies only to the
# dedicated AI account and does not affect any other sudo invocations.
printf 'Defaults>aiops env_keep += "SSH_CONNECTION SSH_CLIENT SSH_ORIGINAL_COMMAND"\n%s ALL=(root) NOPASSWD: %s ""\n' "$AI_USER" "$RUNTIME_COMMAND" >"$sudoers_tmp"
chmod 0440 "$sudoers_tmp"
visudo -cf "$sudoers_tmp" >/dev/null || die "generated sudoers rule is invalid"

cat >"$sshd_tmp" <<SSHD_CONFIG
# Managed by scripts/install-ai-readonly-ssh.sh.
Match User ${AI_USER}
    AuthenticationMethods publickey
    PasswordAuthentication no
    KbdInteractiveAuthentication no
    PubkeyAuthentication yes
    ForceCommand ${RUNTIME_COMMAND}
    DisableForwarding yes
    AllowAgentForwarding no
    AllowTcpForwarding no
    X11Forwarding no
    PermitTTY no
    PermitTunnel no
    PermitUserRC no
    MaxSessions 1

Match all
SSHD_CONFIG
install -o root -g root -m 0644 "$sshd_tmp" "$SSHD_DROPIN"
INSTALLED_SSHD_DROPIN=1
install -d -o root -g root -m 0755 /run/sshd
"$SSHD_BIN" -t || die "sshd rejected the generated configuration"

install -o root -g root -m 0440 "$sudoers_tmp" "$SUDOERS_FILE"
INSTALLED_SUDOERS=1

if ((created_user == 1)); then
  find "$user_home" -mindepth 1 -maxdepth 1 ! -name .ssh -exec rm -rf -- {} +
  rm -rf -- "$user_home/.ssh"
fi
install -d -o root -g root -m 0755 "$user_home" "$user_home/.ssh"
key_options="restrict,command=\"${RUNTIME_COMMAND}\""
if ((ALLOW_ANY_SOURCE == 0)); then
  key_options="from=\"${AI_ALLOWED_FROM}\",${key_options}"
fi
printf '%s %s\n' "$key_options" "$public_key" >"$user_home/.ssh/authorized_keys"
chown root:root "$user_home/.ssh/authorized_keys"
chmod 0644 "$user_home/.ssh/authorized_keys"
INSTALLED_AUTHORIZED_KEYS=1

marker_created_user="$created_user"
if ((created_user == 0)) && ((existing_install == 1)); then
  marker_created_user="$original_created_user"
fi
cat >"$MANAGED_MARKER" <<MARKER
MANAGED_AI_USER=${AI_USER}
CREATED_USER=${marker_created_user}
MARKER
chown root:root "$MANAGED_MARKER"
chmod 0600 "$MANAGED_MARKER"
INSTALLED_MARKER=1

log "Validating effective restrictions"
effective="$("$SSHD_BIN" -T -C "user=${AI_USER},host=localhost,addr=11.57.0.20")"
grep -Fq "forcecommand ${RUNTIME_COMMAND}" <<<"$effective" || die "ForceCommand is not effective for $AI_USER"
grep -Fq 'passwordauthentication no' <<<"$effective" || die "password authentication is not disabled for $AI_USER"
grep -Fq 'permittty no' <<<"$effective" || die "PTY is not disabled for $AI_USER"
grep -Fq 'disableforwarding yes' <<<"$effective" || die "forwarding is not disabled for $AI_USER"

reload_sshd
log "Testing the local forced-command path"
test_output="$(sudo -u "$AI_USER" env SSH_ORIGINAL_COMMAND=help "$RUNTIME_COMMAND")"
grep -Fq 'Allowed commands:' <<<"$test_output" || die "local command-path test failed"

INSTALL_COMMITTED=1
trap - EXIT
rm -f "$key_tmp" "$containers_tmp" "$units_tmp" "$sudoers_tmp" "$sshd_tmp"
if [[ -n "$SNAPSHOT_DIR" ]]; then
  rm -rf -- "$SNAPSHOT_DIR"
  SNAPSHOT_DIR=""
fi

cat <<SUMMARY

Read-only AI SSH access installed.

Account:
  ${AI_USER}

Public key:
  $([[ $ROTATE_PUBLIC_KEY -eq 1 ]] && printf '%s' 'rotated explicitly' || printf '%s' 'installed or unchanged')

Source restriction:
  $([[ $ALLOW_ANY_SOURCE -eq 1 ]] && printf '%s' 'none in authorized_keys' || printf '%s' "$AI_ALLOWED_FROM")

Allowed containers:
$(sed 's/^/  /' "$CONFIG_DIR/containers.allow")

Allowed systemd units:
$(sed 's/^/  /' "$CONFIG_DIR/units.allow")

Test:
  ssh ${AI_USER}@10.77.0.1 help
  ssh ${AI_USER}@10.77.0.1 containers

Uninstall:
  sudo ./scripts/install-ai-readonly-ssh.sh --user ${AI_USER} --uninstall

Keep the private key in the AI backend secret store, never in a prompt,
repository, chat message or VPS filesystem.
SUMMARY