#!/usr/bin/env bash
#
# check_ssh_access.sh
#
# Verifies that SSH (port 22) on a list of VMs is reachable from this machine.
# Run it FROM the dedicated line whose IP is supposed to be allowlisted.
#
# Usage:
#   ./check_ssh_access.sh vms.txt                  # read IPs from a file
#   ./check_ssh_access.sh 10.0.0.5 10.0.0.6        # or pass them inline
#   ./check_ssh_access.sh -p 2222 -t 8 vms.txt     # custom port / timeout
#   ./check_ssh_access.sh -s 103.83.89.157 vms.txt # bind to a specific local IP
#   ./check_ssh_access.sh -n vms.txt               # skip the public-IP check
#   ./check_ssh_access.sh -c vms.txt > report.csv  # CSV output
#
# vms.txt format — one per line, optional label, '#' starts a comment:
#   10.0.0.5     web-01
#   10.0.0.6     db-01
#   vm3.example.com

set -uo pipefail

# ---- your dedicated IPs -----------------------------------------------------
ALLOWED_IPS=("72.255.34.20" "103.83.89.157")

# ---- defaults ---------------------------------------------------------------
PORT=22
TIMEOUT=5
BANNER_WAIT=3
SOURCE_IP=""
SKIP_IP_CHECK=0
CSV=0

usage() { sed -n '3,25p' "$0" | sed 's/^# \{0,1\}//'; exit 0; }

while getopts ":p:t:s:nch" opt; do
  case "$opt" in
    p) PORT="$OPTARG" ;;
    t) TIMEOUT="$OPTARG" ;;
    s) SOURCE_IP="$OPTARG" ;;
    n) SKIP_IP_CHECK=1 ;;
    c) CSV=1 ;;
    h) usage ;;
    *) echo "unknown option -$OPTARG" >&2; exit 2 ;;
  esac
done
shift $((OPTIND - 1))

[ $# -eq 0 ] && { echo "error: give a file of VM IPs or a list of IPs" >&2; exit 2; }

# ---- colours (only on a terminal) -------------------------------------------
if [ -t 1 ] && [ "$CSV" -eq 0 ]; then
  G=$'\e[32m'; R=$'\e[31m'; Y=$'\e[33m'; B=$'\e[1m'; N=$'\e[0m'
else
  G=""; R=""; Y=""; B=""; N=""
fi

# ---- build the target list --------------------------------------------------
TARGETS=()
LABELS=()
add_target() {
  local host="$1" label="${2:-}"
  [ -z "$host" ] && return
  TARGETS+=("$host")
  LABELS+=("${label:-$host}")
}

for arg in "$@"; do
  if [ -f "$arg" ]; then
    while IFS= read -r line || [ -n "$line" ]; do
      line="${line%%#*}"
      line="$(echo "$line" | xargs)"
      [ -z "$line" ] && continue
      add_target $line
    done < "$arg"
  else
    add_target "$arg"
  fi
done

[ ${#TARGETS[@]} -eq 0 ] && { echo "error: no VM IPs found" >&2; exit 2; }

# ---- confirm where we are coming from ---------------------------------------
MY_IP=""
if [ "$SKIP_IP_CHECK" -eq 0 ]; then
  for svc in "https://api.ipify.org" "https://ifconfig.me/ip" "https://icanhazip.com"; do
    MY_IP="$(curl -4 -fsS --max-time 5 "$svc" 2>/dev/null | tr -d '[:space:]')"
    [ -n "$MY_IP" ] && break
  done

  if [ -z "$MY_IP" ]; then
    echo "${Y}warning:${N} could not determine this machine's public IP" >&2
  else
    match=0
    for ip in "${ALLOWED_IPS[@]}"; do
      [ "$MY_IP" = "$ip" ] && match=1
    done
    if [ "$match" -eq 1 ]; then
      [ "$CSV" -eq 0 ] && echo "${B}Source IP:${N} $MY_IP ${G}(dedicated — results are valid)${N}"
    else
      echo "${R}WARNING:${N} your public IP is $MY_IP, which is NOT one of your dedicated IPs" >&2
      echo "         (${ALLOWED_IPS[*]}). Results will not reflect your allowlist." >&2
    fi
  fi
fi

# ---- the probe --------------------------------------------------------------
# Returns:  0 = open (banner on stdout)   1 = refused   2 = filtered/timeout
#           3 = DNS / host resolution failure
probe() {
  local host="$1" out rc

  if [ -n "$SOURCE_IP" ] && command -v nc >/dev/null 2>&1; then
    out="$(nc -w "$TIMEOUT" -s "$SOURCE_IP" "$host" "$PORT" </dev/null 2>/dev/null | head -1)"
    rc=$?
    [ $rc -eq 0 ] && { printf '%s' "$out"; return 0; }
    return 2
  fi

  out="$(timeout "$TIMEOUT" bash -c '
      exec 3<>"/dev/tcp/$0/$1" 2>/dev/null || exit 1
      IFS= read -r -t '"$BANNER_WAIT"' banner <&3
      printf "%s" "$banner"
      exec 3<&-
    ' "$host" "$PORT" 2>/dev/null)"
  rc=$?

  case $rc in
    0)   printf '%s' "$out"; return 0 ;;
    124) return 2 ;;
    *)   getent hosts "$host" >/dev/null 2>&1 || return 3
         return 1 ;;
  esac
}

# ---- run --------------------------------------------------------------------
ok=0; blocked=0; refused=0; bad=0

if [ "$CSV" -eq 1 ]; then
  echo "vm,host,port,status,detail"
else
  printf '\n%-24s %-18s %-11s %s\n' "VM" "HOST" "PORT $PORT" "DETAIL"
  printf '%s\n' "----------------------------------------------------------------------------"
fi

for i in "${!TARGETS[@]}"; do
  host="${TARGETS[$i]}"; label="${LABELS[$i]}"
  banner="$(probe "$host")"; rc=$?

  case $rc in
    0)
      ok=$((ok+1))
      status="ALLOWED"; colour="$G"
      if [ -n "$banner" ]; then
        detail="$(echo "$banner" | tr -d '\r')"
      else
        detail="port open (no banner)"
      fi
      ;;
    1)
      refused=$((refused+1))
      status="REFUSED"; colour="$Y"
      detail="reachable, but connection rejected — sshd down, or a REJECT rule"
      ;;
    2)
      blocked=$((blocked+1))
      status="BLOCKED"; colour="$R"
      detail="no reply within ${TIMEOUT}s — dropped by firewall, or host down"
      ;;
    3)
      bad=$((bad+1))
      status="NO-HOST"; colour="$R"
      detail="cannot resolve hostname"
      ;;
  esac

  if [ "$CSV" -eq 1 ]; then
    printf '%s,%s,%s,%s,"%s"\n' "$label" "$host" "$PORT" "$status" "${detail//\"/\"\"}"
  else
    printf '%-24s %-18s %s%-11s%s %s\n' "$label" "$host" "$colour" "$status" "$N" "$detail"
  fi
done

if [ "$CSV" -eq 0 ]; then
  printf '%s\n' "----------------------------------------------------------------------------"
  printf '%s allowed  |  %s blocked  |  %s refused  |  %s unresolved   (total %s)\n\n' \
    "$ok" "$blocked" "$refused" "$bad" "${#TARGETS[@]}"
fi

[ "$ok" -eq "${#TARGETS[@]}" ] && exit 0
exit 1