#!/bin/bash

set -euo pipefail

# mailcheck.sh — подробная проверка почтового домена (A/MX/servises/SPF/DMARC/DKIM/PTR)
# Работает на macOS и Linux. Требует: dig, host (или nslookup), awk, sed, grep; openssl optional.

PROG_NAME="$(basename "$0")"
DEFAULT_TIMEOUT=5

DEFAULT_SELECTORS=(mail default selector1 s1 google)

usage() {
  cat <<EOF
Usage: $PROG_NAME [options] domain

Options:
  -s, --selector SEL       Проверить только селектор SEL (один).
  -S, --selectors LIST     Список селекторов через запятую (mail,default,...).
  --all-selectors          Перебрать встроенный набор селекторов.
  -t, --timeout N          Таймаут для dig/openssl (в секундах). Default: $DEFAULT_TIMEOUT
  -h, --help               Показать это сообщение.

Примеры:
  $PROG_NAME example.com
  $PROG_NAME -s mail example.com
  $PROG_NAME -S mail,default,s1 example.com
  $PROG_NAME --all-selectors example.com
EOF
  exit 1
}

# Проверки зависимостей
for cmd in dig host awk sed grep; do
  command -v "$cmd" >/dev/null 2>&1 || { echo "ERROR: required command '$cmd' not found"; exit 2; }
done
# openssl — опционально (TLS banner). Если нет, просто пропустим TLS проверку.
OPENSSL_AVAILABLE=1
command -v openssl >/dev/null 2>&1 || OPENSSL_AVAILABLE=0

# парсинг аргументов
SELECTORS=()
ALL_SELECTORS=0
TIMEOUT=$DEFAULT_TIMEOUT

if (( $# == 0 )); then usage; fi

while (( $# )); do
  case "$1" in
    -s|--selector) SELECTORS=("$2"); shift 2 ;;
    -S|--selectors) IFS=',' read -r -a SELECTORS <<< "$2"; shift 2 ;;
    --all-selectors) ALL_SELECTORS=1; shift ;;
    -t|--timeout) TIMEOUT="$2"; shift 2 ;;
    -h|--help) usage ;;
    --) shift; break ;;
    -*)
      echo "Unknown option: $1"; usage ;;
    *) DOMAIN="$1"; shift; break ;;
  esac
done

# если ещё остались аргументы, первый из них — домен
if [[ -z "${DOMAIN:-}" && $# -ge 1 ]]; then DOMAIN="$1"; fi
[[ -z "${DOMAIN:-}" ]] && usage

if (( ALL_SELECTORS == 1 )) && ((${#SELECTORS[@]}==0)); then
  SELECTORS=("${DEFAULT_SELECTORS[@]}")
fi
if ((${#SELECTORS[@]}==0)); then
  # по умолчанию ничего не перебираем — но для DKIM проверим стандартный набор (mail, default, selector1, s1, google)
  SELECTORS=("${DEFAULT_SELECTORS[@]}")
fi

# утилиты-обёртки
dig_short() { dig +time="$TIMEOUT" +tries=2 +short "$@"; }   # возвращает 0..n строк
trim() { sed 's/^[[:space:]]*//; s/[[:space:]]*$//'; }

# ---- helper: resolve A/AAAA to array string (comma separated) ----
resolve_ips() {
  local host="$1"
  local arr=()
  while IFS= read -r ip;  do [[ -n "$ip" ]] && arr+=("$ip"); done < <(dig_short A "$host") || true
  while IFS= read -r ip6; do [[ -n "$ip6" ]] && arr+=("$ip6"); done < <(dig_short AAAA "$host") || true
  if ((${#arr[@]}==0)); then echo "-"; else (IFS=,; echo "${arr[*]}"); fi
}

# ---- helper: PTR for IP ----
ptr_of_ip() {
  local ip="$1"
  if [[ "$ip" =~ : ]]; then
    # IPv6 reverse
    local rev=$(printf "%s" "$ip" | awk -F: '{print $0}' )
    # use dig -x
    dig_short -x "$ip" | sed 's/\.$//' | trim || echo "-"
  else
    dig_short -x "$ip" | sed 's/\.$//' | trim || echo "-"
  fi
}

# ---- helper: check host exists (A/AAAA) ----
host_flag_and_ips() {
  local host="$1"
  local ips
  ips=$(resolve_ips "$host")
  if [[ "$ips" == "-" ]]; then echo "-|-"; else echo "+|$ips"; fi
}

# ---- get MX list -> array of "prio host" ----
get_mx_records() {
  local domain="$1"
  local lines
  lines=$(dig_short MX "$domain" | sort -n || true)
  if [[ -z "$lines" ]]; then
    echo ""
    return
  fi
  # output lines as-is
  echo "$lines"
}

# ---- SPF record (v=spf1) raw ----
get_spf_raw() {
  local domain="$1"
  # Получаем все TXT, убираем кавычки, объединяем строки в одну, нормализуем пробелы
  local txt
  txt="$(dig_short TXT "$domain" | sed 's/^"//; s/"$//' | tr '\n' ' ' | sed 's/[[:space:]]\\+/ /g' | sed 's/^ //; s/ $//')"
  if [[ -z "$txt" ]]; then
    echo "-"
    return
  fi
  # Выбираем часть начиная с v=spf1 (если есть)
  if [[ "$txt" == *v=spf1* ]]; then
    local spf="v=spf1${txt#*v=spf1}"
    # нормализовать пробелы
    spf="$(printf '%s' "$spf" | sed 's/[[:space:]]\\+/ /g' | sed 's/^ //; s/ $//')"
    echo "$spf"
  else
    echo "-"
  fi
}

get_dmarc_raw() {
  local domain="$1"
  local name="_dmarc.$domain"
  local txt
  txt="$(dig_short TXT "$name" | sed 's/^"//; s/"$//' | tr '\n' ' ' | sed 's/[[:space:]]\\+/ /g' | sed 's/^ //; s/ $//')"
  if [[ -z "$txt" ]]; then
    echo "-"
    return
  fi
  # Ищем v=DMARC1 (регистр игнорируем, приводим к верхнему для поиска)
  if echo "$txt" | tr '[:lower:]' '[:upper:]' | grep -q 'V=DMARC1'; then
    # найдём первую позицию v=DMARC1 и вернём от неё до конца строки
    # конструкция: берём оригинальную строку и отрезаем по шаблону v=DMARC1 (без изменения регистра)
    # сначала найдём index в верхнем варианте
    local up="$(printf '%s' "$txt" | tr '[:lower:]' '[:upper:]')"
    local pos
    pos=$(awk -v s="$up" 'BEGIN{print index(s,"V=DMARC1")}')
    if (( pos > 0 )); then
      # отрезаем в оригинальной строке, учитывая смещение
      local tail="${txt:pos-1}"
      tail="$(printf '%s' "$tail" | sed 's/[[:space:]]\\+/ /g' | sed 's/^ //; s/ $//')"
      echo "$tail"
      return
    fi
  fi
  echo "-"
}

# ---- DKIM TXT getter short ----
get_dkim_short() {
  local selector="$1"; local domain="$2"
  local name="${selector}._domainkey.${domain}"
  # try TXT
  local txt
  txt=$(dig_short TXT "$name" | sed 's/^"//; s/"$//' | sed '/^$/d' | tr -d '\r' || true)
  if [[ -n "$txt" ]]; then
    printf '%s' "$txt" | tr -d '\n' | awk '{out=substr($0,1,200); print out (length($0)>200 ? "..." : "")}'; return
  fi
  # try CNAME -> TXT
  local cname
  cname=$(dig_short CNAME "$name" | sed 's/\.$//' | trim || true)
  if [[ -n "$cname" ]]; then
    local txt2
    txt2=$(dig_short TXT "$cname" | sed 's/^"//; s/"$//' | sed '/^$/d' | tr -d '\r' || true)
    if [[ -n "$txt2" ]]; then
      printf '%s' "$txt2" | tr -d '\n' | awk '{out=substr($0,1,200); print out (length($0)>200 ? "..." : "")}'; return
    fi
  fi
  echo "-"
}

# ---- Parse SPF token list for includes, ip4, ip6, a, mx ----
parse_spf_tokens() {
  local spf="$1"
  # tokens separated by spaces
  # we'll print each found mechanism on own line as TYPE|VALUE
  # include:domain, ip4:..., ip6:..., a, mx, ptr (rare)
  echo "$spf" | tr ' ' '\n' | while IFS= read -r tok; do
    case "$tok" in
      include:*) echo "include|${tok#include:}" ;;
      ip4:*) echo "ip4|${tok#ip4:}" ;;
      ip6:*) echo "ip6|${tok#ip6:}" ;;
      a) echo "a|$2" ;;           # a of domain
      a:* ) echo "a|${tok#a:}" ;;
      mx) echo "mx|$2" ;;
      mx:* ) echo "mx|${tok#mx:}" ;;
      ptr) echo "ptr|$tok" ;;
      redirect=*) echo "redirect|${tok#redirect=}" ;;
      exp=* ) echo "exp|${tok#exp=}" ;;
      ~all|+all|all|?all|-all) echo "all|${tok}" ;;
      *) ;; # ignore other modifiers (~, -)
    esac
  done
}

# ---- Print header nice ----
print_header() {
  echo
  echo "========================================================"
  echo " Mail/Domain quickcheck: $DOMAIN"
  echo " $(date)"
  echo "--------------------------------------------------------"
}

# ---- main output ----
print_header

# 1) A / AAAA records for domain (root)
echo
echo "-- A / AAAA for $DOMAIN:"
ips_root=$(resolve_ips "$DOMAIN")
echo "  $DOMAIN -> $ips_root"

# 2) MX records
echo
echo "-- MX records for $DOMAIN:"
mx_raw=$(get_mx_records "$DOMAIN")
if [[ -z "$mx_raw" ]]; then
  echo "  No MX records."
else
  # iterate
  printf "  %-5s %-40s %s\n" "PRIO" "MX_HOST" "IP(s) [PTR]"
  while IFS= read -r mxline; do
    [[ -z "$mxline" ]] && continue
    prio=$(awk '{print $1}' <<<"$mxline")
    host=$(awk '{print $2}' <<<"$mxline" | sed 's/\.$//')
    ips=$(resolve_ips "$host")
    # for each ip, get PTR (show only first 2)
    ptrs=""
    if [[ "$ips" != "-" ]]; then
      IFS=',' read -r -a iparr <<< "$ips"
      local_ptrs=()
      for ip in "${iparr[@]:0:3}"; do
        p=$(ptr_of_ip "$ip")
        local_ptrs+=("${p:--}")
      done
      ptrs=$(IFS=,; echo "${local_ptrs[*]}")
    else
      ptrs="-"
    fi
    printf "  %-5s %-40s %s [PTR: %s]\n" "$prio" "$host" "$ips" "$ptrs"
  done <<< "$mx_raw"
fi

# 3) Mail/IMAP/POP/SMTP host checks
echo
echo "-- Common service hosts (mail/imap/pop/smtp) for $DOMAIN:"
printf "  %-20s %-4s %-20s %s\n" "HOST" "OK" "IP(s)" "PTR(example)"
for svc in mail imap pop pop3 smtp submission smtps; do
  host="${svc}.${DOMAIN}"
  res=$(host_flag_and_ips "$host")
  flag="${res%%|*}"
  ips="${res#*|}"
  ptr_example="-"
  if [[ "$ips" != "-" ]]; then
    IFS=',' read -r -a iparr <<< "$ips"
    ptr_example="$(ptr_of_ip "${iparr[0]}")"
  fi
  printf "  %-20s %-4s %-20s %s\n" "$host" "$flag" "$ips" "$ptr_example"
done

# 4) SPF
echo
echo "-- SPF (v=spf1) for $DOMAIN:"
spf_raw=$(get_spf_raw "$DOMAIN")
if [[ "$spf_raw" == "-" ]]; then
  echo "  No SPF record (v=spf1) found."
else
  echo "  $spf_raw"
  echo "  Parsed mechanisms:"
  parse_spf_tokens "$spf_raw" "$DOMAIN" | while IFS='|' read -r typ val; do
    case "$typ" in
      include)
        echo "    include: $val"
        # try to fetch spf of included host quickly
        incspf=$(get_spf_raw "$val")
        echo "      -> included domain SPF: ${incspf:--}"
        ;;
      ip4|ip6)
        echo "    $typ: $val"
        ;;
      a)
        echo "    a: resolves to $(resolve_ips "$val")"
        ;;
      mx)
        echo "    mx: MX hosts resolve to:"
        # expand mx of that domain
        mxs=$(get_mx_records "$val")
        if [[ -z "$mxs" ]]; then echo "       (no MX)"; else
          while IFS= read -r L; do
            [[ -z "$L" ]] && continue
            h=$(awk '{print $2}' <<<"$L" | sed 's/\.$//')
            echo "       - $h -> $(resolve_ips "$h")"
          done <<< "$mxs"
        fi
        ;;
      redirect)
        echo "    redirect -> $val"
        ;;
      all)
        echo "    all mechanism: $val"
        ;;
      *)
        echo "    $typ: $val"
        ;;
    esac
  done
fi

# 5) DMARC
echo
echo "-- DMARC (_dmarc.$DOMAIN) :"
dmarc_raw=$(get_dmarc_raw "$DOMAIN")
if [[ "$dmarc_raw" == "-" ]]; then
  echo "  No DMARC record found."
else
  echo "  $dmarc_raw"
  # parse basic tags: p=, rua=, ruf=, pct=
  p=$(echo "$dmarc_raw" | sed -n 's/.*p=\([a-z]*\).*/\1/p' || true)
  rua=$(echo "$dmarc_raw" | sed -n 's/.*rua=\([^;]*\).*/\1/p' || true)
  ruf=$(echo "$dmarc_raw" | sed -n 's/.*ruf=\([^;]*\).*/\1/p' || true)
  pct=$(echo "$dmarc_raw" | sed -n 's/.*pct=\([0-9]*\).*/\1/p' || true)
  echo "    policy: ${p:-(not set)}"
  echo "    rua: ${rua:-(not set)}"
  echo "    ruf: ${ruf:-(not set)}"
  echo "    pct: ${pct:-(not set)}"
fi

# 6) DKIM checks: iterate selectors
echo
echo "-- DKIM selectors check for $DOMAIN:"
for sel in "${SELECTORS[@]}"; do
  val=$(get_dkim_short "$sel" "$DOMAIN")
  printf "  %-12s %s\n" "$sel" "$val"
done

# 7) Additional PTR checks: evaluate MX IPs PTR->A (forward-confirmed reverse DNS)
echo
echo "-- PTR / forward-confirmation for MX IPs:"
if [[ -n "$mx_raw" ]]; then
  while IFS= read -r mxline; do
    [[ -z "$mxline" ]] && continue
    host=$(awk '{print $2}' <<<"$mxline" | sed 's/\.$//')
    ips=$(resolve_ips "$host")
    if [[ "$ips" == "-" ]]; then
      echo "  MX $host -> no A/AAAA"
      continue
    fi
    IFS=',' read -r -a iparr <<< "$ips"
    for ip in "${iparr[@]}"; do
      p=$(ptr_of_ip "$ip")
      # forward lookup of PTR
      fwd="-"
      if [[ "$p" != "-" ]]; then
        fwd=$(resolve_ips "$p")
      fi
      printf "  %-40s %-20s PTR=%-40s FORWARD=%s\n" "$host" "$ip" "$p" "$fwd"
    done
  done <<< "$mx_raw"
else
  echo "  no MX => skip PTR checks"
fi

# 8) Extra: check SPF includes' referenced IPs (quick)
echo
echo "-- SPF includes expanded (quick):"
if [[ "$spf_raw" != "-" ]]; then
  echo "$spf_raw" | tr ' ' '\n' | grep -E '^(include:|ip4:|ip6:|a|mx:)' | while IFS= read -r L; do
    case "$L" in
      include:*)
        inc="${L#include:}"
        echo "  include: $inc -> SPF: $(get_spf_raw "$inc")"
        ;;
      ip4:*|ip6:*)
        echo "  $L"
        ;;
      a)
        echo "  a -> $(resolve_ips "$DOMAIN")"
        ;;
      mx|mx:*)
        echo "  $L"
        ;;
      *)
        echo "  $L"
        ;;
    esac
  done
else
  echo "  no spf to expand"
fi

# 9) Optional: test SMTP TLS connect to first MX (openssl) — если openssl доступен
if (( OPENSSL_AVAILABLE )); then
  echo
  echo "-- SMTP STARTTLS/TLS probe to first MX (if reachable):"
  first_mx_host=""
  if [[ -n "$mx_raw" ]]; then
    first_mx_host=$(awk 'NR==1{print $2}' <<<"$mx_raw" | sed 's/\.$//')
  fi
  if [[ -n "$first_mx_host" ]]; then
    mx_ips=$(resolve_ips "$first_mx_host")
    if [[ "$mx_ips" == "-" ]]; then
      echo "  first MX $first_mx_host has no A/AAAA"
    else
      # pick first IP
      IFS=',' read -r first_ip _ <<< "$mx_ips"
      echo "  probing $first_mx_host ($first_ip) :"
      # try plain SMTP port 25 banner and EHLO + STARTTLS
      echo -n "    SMTP banner: "
      timeout "$TIMEOUT" sh -c "echo | nc -w $TIMEOUT $first_ip 25 2>/dev/null" || true
      # Try openssl s_client to port 25 STARTTLS (may fail if server doesn't allow plain TLS)
      echo "    trying openssl s_client -starttls smtp -crlf -connect ${first_ip}:25 (timeout $TIMEOUT s)"
      if timeout "$TIMEOUT" openssl s_client -starttls smtp -crlf -connect "${first_ip}:25" </dev/null 2>&1 | sed -n '1,20p'; then
        echo "    (done)"
      else
        echo "    (openssl STARTTLS probe failed or timed out)"
      fi
    fi
  else
    echo "  no MX to probe"
  fi
else
  echo
  echo "-- openssl not available, skipping SMTP TLS probe"
fi

echo
echo "========================================================"
echo "Done."

