#!/usr/bin/env bash
set -u

MODE="check"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TARGET_HOST="${WINDOWS99_HOST:-192.168.0.99}"
CONNECT_TIMEOUT="${WINDOWS99_CONNECT_TIMEOUT:-3}"
SSH_TIMEOUT="${WINDOWS99_SSH_TIMEOUT:-3}"
REMOTE_LOCATOR_TIMEOUT="${WINDOWS99_REMOTE_LOCATOR_TIMEOUT:-60}"
SSH_PORT="${WINDOWS99_SSH_PORT:-22}"
VIA_HOST="${WINDOWS99_VIA_HOST:-}"
VIA_HOST_PORT="${WINDOWS99_VIA_PORT:-22}"
MAX_AUTH_USERS="${WINDOWS99_MAX_AUTH_USERS:-}"
MAX_AUTH_USERS_EXPLICIT=0
if [[ -n "${WINDOWS99_MAX_AUTH_USERS:-}" ]]; then
  MAX_AUTH_USERS_EXPLICIT=1
fi
KNOWN_HOSTS_FILE="${WINDOWS99_KNOWN_HOSTS_FILE:-/tmp/awoooi-windows99-known_hosts}"
LOCAL_LOCATOR_SCRIPT="${WINDOWS99_LOCAL_LOCATOR_SCRIPT:-${SCRIPT_DIR}/windows99-vmx-source-locator.ps1}"
REMOTE_LOCATOR_COMMAND="${WINDOWS99_REMOTE_LOCATOR_COMMAND:-powershell -NoProfile -ExecutionPolicy Bypass -Command \"& ([scriptblock]::Create([Console]::In.ReadToEnd())) -Mode Check\"}"
SSH_USERS=(ogt wooo ooo administrator Administrator)
SSH_USERS_EXPLICIT=0

if [[ -n "${WINDOWS99_SSH_USERS:-}" ]]; then
  # shellcheck disable=SC2206
  SSH_USERS=(${WINDOWS99_SSH_USERS})
  SSH_USERS_EXPLICIT=1
fi

is_positive_int() {
  [[ "$1" =~ ^[1-9][0-9]*$ ]]
}

usage() {
  printf '%s\n' "usage: $0 [--check|--collect] [--host HOST] [--users 'u1 u2'] [--via-host USER@HOST] [--timeout SECONDS] [--max-auth-users N]"
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --check)
      MODE="check"
      ;;
    --collect)
      MODE="collect"
      ;;
    --host)
      shift
      TARGET_HOST="${1:-}"
      ;;
    --users)
      shift
      # shellcheck disable=SC2206
      SSH_USERS=(${1:-})
      SSH_USERS_EXPLICIT=1
      ;;
    --via-host)
      shift
      VIA_HOST="${1:-}"
      ;;
    --via-port)
      shift
      VIA_HOST_PORT="${1:-22}"
      ;;
    --timeout)
      shift
      CONNECT_TIMEOUT="${1:-5}"
      SSH_TIMEOUT="${CONNECT_TIMEOUT}"
      ;;
    --max-auth-users)
      shift
      MAX_AUTH_USERS="${1:-}"
      MAX_AUTH_USERS_EXPLICIT=1
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      printf '%s\n' "error=unknown_argument:$1" >&2
      usage >&2
      exit 64
      ;;
  esac
  shift
done

if ! is_positive_int "${CONNECT_TIMEOUT}"; then
  CONNECT_TIMEOUT=3
fi
if ! is_positive_int "${SSH_TIMEOUT}"; then
  SSH_TIMEOUT=3
fi
if ! is_positive_int "${REMOTE_LOCATOR_TIMEOUT}"; then
  REMOTE_LOCATOR_TIMEOUT=60
fi
if ! is_positive_int "${MAX_AUTH_USERS}"; then
  MAX_AUTH_USERS="${#SSH_USERS[@]}"
fi
if [[ "${SSH_USERS_EXPLICIT}" == "1" && "${MAX_AUTH_USERS_EXPLICIT}" != "1" ]]; then
  MAX_AUTH_USERS="${#SSH_USERS[@]}"
fi
if ! is_positive_int "${MAX_AUTH_USERS}"; then
  MAX_AUTH_USERS=2
fi

if [[ "${MODE}" != "check" && "${MODE}" != "collect" ]]; then
  printf '%s\n' "error=invalid_mode:${MODE}" >&2
  exit 64
fi

PORT_TIMEOUT_WRAPPER="none"
if command -v timeout >/dev/null 2>&1; then
  PORT_TIMEOUT_WRAPPER="timeout"
elif command -v gtimeout >/dev/null 2>&1; then
  PORT_TIMEOUT_WRAPPER="gtimeout"
fi

port_open() {
  local port="$1"
  if ! command -v nc >/dev/null 2>&1; then
    return 1
  fi
  if [[ "${PORT_TIMEOUT_WRAPPER}" == "timeout" ]]; then
    timeout "$((CONNECT_TIMEOUT + 1))s" nc -z -w "${CONNECT_TIMEOUT}" "${TARGET_HOST}" "${port}" >/dev/null 2>&1
  elif [[ "${PORT_TIMEOUT_WRAPPER}" == "gtimeout" ]]; then
    gtimeout "$((CONNECT_TIMEOUT + 1))s" nc -z -w "${CONNECT_TIMEOUT}" "${TARGET_HOST}" "${port}" >/dev/null 2>&1
  else
    nc -z -w "${CONNECT_TIMEOUT}" "${TARGET_HOST}" "${port}" >/dev/null 2>&1
  fi
}

bool_for_port() {
  local port="$1"
  if port_open "${port}"; then
    printf '1'
  else
    printf '0'
  fi
}

join_users() {
  local joined=""
  local user
  for user in "${SSH_USERS[@]}"; do
    if [[ -z "${joined}" ]]; then
      joined="${user}"
    else
      joined="${joined},${user}"
    fi
  done
  printf '%s' "${joined}"
}

PORT_22_OPEN="$(bool_for_port 22)"
PORT_3389_OPEN="$(bool_for_port 3389)"
PORT_5985_OPEN="$(bool_for_port 5985)"
PORT_5986_OPEN="$(bool_for_port 5986)"

SSH_BATCHMODE_AUTH_READY=0
SSH_AUTHENTICATED_USER=""
SSH_AUTH_PROBE_EXIT_STATUS="not_attempted"
SSH_AUTH_PROBE_STDOUT_PRESENT=0
SSH_AUTH_ATTEMPTED_USERS="$(join_users)"
SSH_AUTH_PROBED_USERS=0
SSH_TIMEOUT_WRAPPER="none"
if command -v timeout >/dev/null 2>&1; then
  SSH_TIMEOUT_WRAPPER="timeout"
elif command -v gtimeout >/dev/null 2>&1; then
  SSH_TIMEOUT_WRAPPER="gtimeout"
fi

SSH_OPTS=(
  -o BatchMode=yes
  -o PreferredAuthentications=publickey
  -o PubkeyAuthentication=yes
  -o PasswordAuthentication=no
  -o KbdInteractiveAuthentication=no
  -o NumberOfPasswordPrompts=0
  -o ConnectTimeout="${SSH_TIMEOUT}"
  -o ConnectionAttempts=1
  -o GSSAPIAuthentication=no
  -o LogLevel=ERROR
  -o StrictHostKeyChecking=no
  -o UserKnownHostsFile="${KNOWN_HOSTS_FILE}"
  -p "${SSH_PORT}"
)

VIA_SSH_OPTS=(
  -o BatchMode=yes
  -o PreferredAuthentications=publickey
  -o PubkeyAuthentication=yes
  -o PasswordAuthentication=no
  -o KbdInteractiveAuthentication=no
  -o NumberOfPasswordPrompts=0
  -o ConnectTimeout="${SSH_TIMEOUT}"
  -o ConnectionAttempts=1
  -o GSSAPIAuthentication=no
  -o LogLevel=ERROR
  -o StrictHostKeyChecking=no
  -o UserKnownHostsFile="${KNOWN_HOSTS_FILE}"
  -p "${VIA_HOST_PORT}"
)

quote_words() {
  printf '%q ' "$@"
}

run_ssh_timed() {
  local user="$1"
  local timeout_seconds="$2"
  local timeout_wrapper_seconds="$((timeout_seconds + 1))"
  local inner_command=""
  shift 2
  if ! is_positive_int "${timeout_seconds}"; then
    timeout_seconds="${SSH_TIMEOUT}"
    timeout_wrapper_seconds="$((SSH_TIMEOUT + 1))"
  fi
  if [[ -n "${VIA_HOST}" ]]; then
    local inner_ssh=(ssh "${SSH_OPTS[@]}" -o ConnectTimeout="${timeout_seconds}" "${user}@${TARGET_HOST}" "$@")
    inner_command="$(quote_words "${inner_ssh[@]}")"
    if [[ "${SSH_TIMEOUT_WRAPPER}" == "timeout" ]]; then
      timeout "${timeout_wrapper_seconds}s" ssh "${VIA_SSH_OPTS[@]}" "${VIA_HOST}" "bash -lc $(printf '%q' "${inner_command}")"
    elif [[ "${SSH_TIMEOUT_WRAPPER}" == "gtimeout" ]]; then
      gtimeout "${timeout_wrapper_seconds}s" ssh "${VIA_SSH_OPTS[@]}" "${VIA_HOST}" "bash -lc $(printf '%q' "${inner_command}")"
    else
      ssh "${VIA_SSH_OPTS[@]}" "${VIA_HOST}" "bash -lc $(printf '%q' "${inner_command}")"
    fi
    return $?
  fi
  if [[ "${SSH_TIMEOUT_WRAPPER}" == "timeout" ]]; then
    timeout "${timeout_wrapper_seconds}s" ssh "${SSH_OPTS[@]}" -o ConnectTimeout="${timeout_seconds}" "${user}@${TARGET_HOST}" "$@"
  elif [[ "${SSH_TIMEOUT_WRAPPER}" == "gtimeout" ]]; then
    gtimeout "${timeout_wrapper_seconds}s" ssh "${SSH_OPTS[@]}" -o ConnectTimeout="${timeout_seconds}" "${user}@${TARGET_HOST}" "$@"
  else
    ssh "${SSH_OPTS[@]}" -o ConnectTimeout="${timeout_seconds}" "${user}@${TARGET_HOST}" "$@"
  fi
}

run_ssh() {
  local user="$1"
  shift
  run_ssh_timed "${user}" "${SSH_TIMEOUT}" "$@"
}

if [[ "${PORT_22_OPEN}" == "1" ]]; then
  for user in "${SSH_USERS[@]}"; do
    if [[ "${SSH_AUTH_PROBED_USERS}" -ge "${MAX_AUTH_USERS}" ]]; then
      break
    fi
    SSH_AUTH_PROBED_USERS=$((SSH_AUTH_PROBED_USERS + 1))
    auth_output=""
    if auth_output="$(run_ssh "${user}" "echo AWOOOI_WINDOWS99_SSH_READY" 2>&1)"; then
      SSH_BATCHMODE_AUTH_READY=1
      SSH_AUTHENTICATED_USER="${user}"
      SSH_AUTH_PROBE_EXIT_STATUS=0
      if [[ -n "${auth_output}" ]]; then
        SSH_AUTH_PROBE_STDOUT_PRESENT=1
      fi
      break
    else
      SSH_AUTH_PROBE_EXIT_STATUS=$?
    fi
  done
fi

DRY_RUN="true"
REMOTE_LOCATOR_ATTEMPTED=0
REMOTE_LOCATOR_EXIT_STATUS="not_attempted"
LOCATOR_COLLECTION_STATUS="blocked_ssh_publickey_auth_missing"
SAFE_NEXT_STEP="select_existing_authorized_public_key_user_or_set_WINDOWS99_SSH_USERS_then_rerun_locator_no_password"
REMOTE_LOCATOR_OUTPUT=""
LOCAL_LOCATOR_SCRIPT_PRESENT=0
if [[ -f "${LOCAL_LOCATOR_SCRIPT}" ]]; then
  LOCAL_LOCATOR_SCRIPT_PRESENT=1
fi
PROCESS_EXIT_STATUS=0

if [[ "${PORT_22_OPEN}" != "1" ]]; then
  LOCATOR_COLLECTION_STATUS="blocked_ssh_port_closed"
  SAFE_NEXT_STEP="enable_existing_ssh_management_channel_publickey_only_then_rerun_locator_no_secret"
  if [[ "${MODE}" == "collect" ]]; then
    PROCESS_EXIT_STATUS=75
  fi
elif [[ "${SSH_BATCHMODE_AUTH_READY}" != "1" ]]; then
  LOCATOR_COLLECTION_STATUS="blocked_ssh_publickey_auth_missing"
  SAFE_NEXT_STEP="select_existing_authorized_public_key_user_or_set_WINDOWS99_SSH_USERS_then_rerun_locator_no_password"
  if [[ "${MODE}" == "collect" ]]; then
    PROCESS_EXIT_STATUS=75
  fi
elif [[ "${MODE}" == "check" ]]; then
  LOCATOR_COLLECTION_STATUS="ready_ssh_batchmode_auth_probe_only"
  SAFE_NEXT_STEP="rerun_locator_with_collect_then_commit_no_secret_locator_artifact_and_scorecard_rerun"
elif [[ "${LOCAL_LOCATOR_SCRIPT_PRESENT}" != "1" ]]; then
  DRY_RUN="false"
  LOCATOR_COLLECTION_STATUS="blocked_local_locator_script_missing"
  SAFE_NEXT_STEP="restore_windows99_vmx_source_locator_ps1_source_then_rerun_no_secret_no_remote_write"
  PROCESS_EXIT_STATUS=75
else
  DRY_RUN="false"
  REMOTE_LOCATOR_ATTEMPTED=1
  if REMOTE_LOCATOR_OUTPUT="$(run_ssh_timed "${SSH_AUTHENTICATED_USER}" "${REMOTE_LOCATOR_TIMEOUT}" "${REMOTE_LOCATOR_COMMAND}" <"${LOCAL_LOCATOR_SCRIPT}" 2>&1)"; then
    REMOTE_LOCATOR_OUTPUT="${REMOTE_LOCATOR_OUTPUT//$'\r'/}"
    REMOTE_LOCATOR_EXIT_STATUS=0
    if grep -q '^AWOOOI_WINDOWS99_VMX_SOURCE_LOCATOR=1$' <<<"${REMOTE_LOCATOR_OUTPUT}" && grep -q '^MODE=Check$' <<<"${REMOTE_LOCATOR_OUTPUT}"; then
      LOCATOR_COLLECTION_STATUS="collected_windows99_vmx_source_locator_stdout"
      SAFE_NEXT_STEP="commit_no_secret_locator_artifact_then_rerun_reboot_auto_recovery_slo_scorecard"
    else
      LOCATOR_COLLECTION_STATUS="blocked_remote_locator_output_invalid"
      SAFE_NEXT_STEP="inspect_no_secret_locator_stdout_then_fix_in_memory_locator_script_and_rerun"
      PROCESS_EXIT_STATUS=75
    fi
  else
    REMOTE_LOCATOR_EXIT_STATUS=$?
    LOCATOR_COLLECTION_STATUS="blocked_remote_locator_command_failed"
    SAFE_NEXT_STEP="inspect_no_secret_locator_stdout_then_fix_in_memory_locator_script_and_rerun"
    PROCESS_EXIT_STATUS=75
  fi
fi

printf '%s\n' "schema_version=windows99_vmx_source_locator_collector_v1"
printf '%s\n' "dry_run=${DRY_RUN}"
printf '%s\n' "target_host=${TARGET_HOST}"
printf '%s\n' "via_host=${VIA_HOST:-none}"
printf '%s\n' "target_host_alias=99"
printf '%s\n' "target_vm_alias=111"
printf '%s\n' "connect_timeout_seconds=${CONNECT_TIMEOUT}"
printf '%s\n' "ssh_timeout_seconds=${SSH_TIMEOUT}"
printf '%s\n' "remote_locator_timeout_seconds=${REMOTE_LOCATOR_TIMEOUT}"
printf '%s\n' "port_timeout_wrapper=${PORT_TIMEOUT_WRAPPER}"
printf '%s\n' "ssh_auth_probe_user_limit=${MAX_AUTH_USERS}"
printf '%s\n' "ssh_timeout_wrapper=${SSH_TIMEOUT_WRAPPER}"
printf '%s\n' "port_22_open=${PORT_22_OPEN}"
printf '%s\n' "port_3389_open=${PORT_3389_OPEN}"
printf '%s\n' "port_5985_open=${PORT_5985_OPEN}"
printf '%s\n' "port_5986_open=${PORT_5986_OPEN}"
printf '%s\n' "ssh_candidate_users=${SSH_AUTH_ATTEMPTED_USERS}"
printf '%s\n' "ssh_auth_probed_users=${SSH_AUTH_PROBED_USERS}"
printf '%s\n' "ssh_batchmode_auth_ready=${SSH_BATCHMODE_AUTH_READY}"
printf '%s\n' "ssh_authenticated_user=${SSH_AUTHENTICATED_USER}"
printf '%s\n' "ssh_auth_probe_exit_status=${SSH_AUTH_PROBE_EXIT_STATUS}"
printf '%s\n' "ssh_auth_probe_stdout_present=${SSH_AUTH_PROBE_STDOUT_PRESENT}"
printf '%s\n' "remote_locator_mode=in_memory_stdin_scriptblock"
printf '%s\n' "local_locator_script_present=${LOCAL_LOCATOR_SCRIPT_PRESENT}"
printf '%s\n' "remote_locator_attempted=${REMOTE_LOCATOR_ATTEMPTED}"
printf '%s\n' "remote_locator_exit_status=${REMOTE_LOCATOR_EXIT_STATUS}"
printf '%s\n' "locator_collection_status=${LOCATOR_COLLECTION_STATUS}"
printf '%s\n' "safe_next_step=${SAFE_NEXT_STEP}"
printf '%s\n' "secret_value_read=false"
printf '%s\n' "password_prompt_allowed=false"
printf '%s\n' "remote_write_performed=false"
printf '%s\n' "host_reboot_performed=false"
printf '%s\n' "vm_power_change_performed=false"
printf '%s\n' "windows_registry_apply_performed=false"
printf '%s\n' "scheduled_task_write_performed=false"
printf '%s\n' "vmx_file_content_read=false"
printf '%s\n' "raw_path_output=false"

if [[ "${REMOTE_LOCATOR_ATTEMPTED}" == "1" ]]; then
  printf '%s\n' "remote_locator_output_begin"
  printf '%s\n' "${REMOTE_LOCATOR_OUTPUT}"
  printf '%s\n' "remote_locator_output_end"
fi

exit "${PROCESS_EXIT_STATUS}"