#!/bin/bash
set -euo pipefail

# Quick fail2ban setup with a strict risk profile (Debian/Ubuntu).
# Strict defaults: 3 failed attempts -> 24h ban; repeat offenders -> 1 week.
# Override via env, e.g. SSH_PORT=2222 SSHD_MODE=extra ./setup-fail2ban.sh
#
# Notes on the defaults (verified against fail2ban 0.11/1.0/1.1 and the Debian
# and Ubuntu packages):
#   * The systemd backend is used for sshd so the jail works on hosts without
#     /var/log/auth.log (i.e. without rsyslog). It requires python3-systemd,
#     which Debian/Ubuntu install with fail2ban.
#   * A systemd backend makes fail2ban ignore "logpath" and use the filter's
#     journalmatch instead - so the [recidive] jail, whose log only exists as a
#     file, must keep a file backend, otherwise it silently never bans.
#   * SSHD_MODE defaults to "normal": "aggressive" also counts pre-auth client
#     disconnects ("Connection closed by authenticating user ... [preauth]") as
#     failures, which bans legitimate users with a 24h/1w profile.
#   * "dbpurgeage" is a fail2ban.conf [Definition] option; setting it here would
#     be silently ignored, so it is not written.
#   * The effective logtarget (fail2ban.conf, then fail2ban.d/*.conf, then
#     fail2ban.local) must resolve to a file: [recidive] reads it to count bans,
#     and a SYSLOG/STDERR target would leave that jail silently inert. The value
#     is therefore resolved, validated and verified to actually receive log
#     output before the script reports success.

SSH_PORT="${SSH_PORT:-}"
IGNORE_IP="${IGNORE_IP:-}"
SSHD_MODE="${SSHD_MODE:-normal}"
MAXRETRY="${MAXRETRY:-3}"
FINDTIME="${FINDTIME:-10m}"
BANTIME="${BANTIME:-24h}"
RECIDIVE_MAXRETRY="${RECIDIVE_MAXRETRY:-3}"
RECIDIVE_FINDTIME="${RECIDIVE_FINDTIME:-1d}"
RECIDIVE_BANTIME="${RECIDIVE_BANTIME:-1w}"

JAIL_LOCAL="/etc/fail2ban/jail.local"
JAIL_D="/etc/fail2ban/jail.d"

if [ "$EUID" -ne 0 ]; then
  echo "Please run as root (sudo)." >&2
  exit 1
fi

if [ ! -r /etc/os-release ]; then
  echo "/etc/os-release is missing; cannot detect the distribution." >&2
  exit 1
fi

. /etc/os-release

case "${ID:-}" in
  debian|ubuntu|raspbian) ;;
  *)
    echo "Unsupported distribution: ${ID:-unknown} (only Debian, Ubuntu and Raspbian are supported)." >&2
    exit 1
    ;;
esac

# fail2ban silently ignores an unknown filter mode, which would leave sshd
# protected by less (or more) than the operator expects.
case "$SSHD_MODE" in
  normal|ddos|extra|aggressive) ;;
  *)
    echo "Invalid SSHD_MODE: '$SSHD_MODE' (expected normal, ddos, extra or aggressive)." >&2
    exit 1
    ;;
esac

export DEBIAN_FRONTEND=noninteractive

sshd_effective() {
  local key="$1" out=""
  out="$(timeout 10 sshd -T 2>/dev/null || true)"
  if [ -n "$out" ]; then
    printf '%s\n' "$out" | awk -v k="$key" '$1 == k { print $2 }'
    return 0
  fi
  awk -v k="$key" 'tolower($1) == k { print $2 }' /etc/ssh/sshd_config 2>/dev/null || true
}

detect_ssh_port() {
  local ports
  ports="$(sshd_effective port | sort -u | paste -sd, -)"
  if [ -z "$ports" ]; then
    ports="$(ss -H -tlnp 2>/dev/null | awk '/sshd/ { n = split($4, a, ":"); print a[n] }' | sort -u | paste -sd, -)"
  fi
  printf '%s' "${ports:-22}"
}

detect_client_ip() {
  local ip="" tty_name=""
  if [ -n "${SSH_CONNECTION:-}" ]; then
    ip="${SSH_CONNECTION%% *}"
  fi
  if [ -z "$ip" ] && [ -n "${SSH_CLIENT:-}" ]; then
    ip="${SSH_CLIENT%% *}"
  fi
  if [ -z "$ip" ]; then
    # sudo's env_reset usually drops SSH_CONNECTION/SSH_CLIENT; fall back to the
    # utmp entry of our own terminal. Note that "who am i" cannot be used here:
    # it resolves the entry through stdin, which is a pipe for the documented
    # "curl ... | sudo bash" invocation.
    tty_name="$(tty 2>/dev/null || true)"
    case "$tty_name" in
      /dev/*)
        ip="$(who 2>/dev/null | awk -v t="${tty_name#/dev/}" '$2 == t { print $NF }' | tail -n1 | tr -d '()' || true)"
        ;;
    esac
  fi
  case "$ip" in
    *[!0-9a-fA-F:.]*|"") ip="" ;;
  esac
  printf '%s' "$ip"
}

detect_banaction() {
  if command -v nft >/dev/null 2>&1; then
    printf 'nftables-multiport'
  elif command -v iptables >/dev/null 2>&1; then
    printf 'iptables-multiport'
  else
    echo "Neither nft nor iptables is available; fail2ban cannot enforce bans." >&2
    exit 1
  fi
}

# [recidive] (and other "allports" jails) ignore banaction and use
# banaction_allports, whose default on older Debian/Ubuntu is iptables-allports.
# Keep both in sync with the firewall actually present.
detect_banaction_allports() {
  case "$1" in
    nftables-*) printf 'nftables-allports' ;;
    iptables-*) printf 'iptables-allports' ;;
    *)
      echo "Cannot derive an all-ports banaction from '$1'." >&2
      exit 1
      ;;
  esac
}

# fail2ban reads fail2ban.conf, then fail2ban.d/*.conf, then fail2ban.local, so
# the value from the last file that sets logtarget is the effective one.
logtarget_effective() {
  local f value="" current
  for f in /etc/fail2ban/fail2ban.conf /etc/fail2ban/fail2ban.d/*.conf \
           /etc/fail2ban/fail2ban.local; do
    [ -r "$f" ] || continue
    current="$(awk '
      /^[[:space:]]*logtarget[[:space:]]*=/ {
        line = $0
        sub(/^[^=]*=[[:space:]]*/, "", line)
        sub(/[[:space:]]*#.*/, "", line)
        gsub(/^[[:space:]]+|[[:space:]]+$/, "", line)
        value = line
      }
      END { if (value != "") print value }
    ' "$f")"
    if [ -n "$current" ]; then
      value="$current"
    fi
  done
  printf '%s' "$value"
}

echo "==> Installing fail2ban"
apt-get update
apt-get install -y fail2ban

echo "==> Resolving fail2ban log target"
FAIL2BAN_LOG="$(logtarget_effective)"
case "$FAIL2BAN_LOG" in
  \"*\") FAIL2BAN_LOG="${FAIL2BAN_LOG#\"}"; FAIL2BAN_LOG="${FAIL2BAN_LOG%\"}" ;;
  \'*\') FAIL2BAN_LOG="${FAIL2BAN_LOG#\'}"; FAIL2BAN_LOG="${FAIL2BAN_LOG%\'}" ;;
esac

# [recidive] can only count bans in a file. Refuse to configure a jail that
# could never fire instead of leaving the host with a false sense of coverage.
case "$FAIL2BAN_LOG" in
  '')
    echo "Cannot determine fail2ban's logtarget: no logtarget in fail2ban.conf," >&2
    echo "fail2ban.d/*.conf or fail2ban.local. The [recidive] jail reads that log." >&2
    exit 1
    ;;
  /*) ;;
  *)
    echo "fail2ban logs to '$FAIL2BAN_LOG', which is not a file path." >&2
    echo "The [recidive] jail reads the log file and would never see a ban." >&2
    echo "Set 'logtarget = /var/log/fail2ban.log' in /etc/fail2ban/fail2ban.local" >&2
    echo "and run this script again." >&2
    exit 1
    ;;
esac

case "$FAIL2BAN_LOG" in
  *[!A-Za-z0-9._/-]*)
    echo "Refusing suspicious logtarget path: '$FAIL2BAN_LOG'." >&2
    exit 1
    ;;
esac

echo "==> Detecting SSH and firewall"
SSH_PORT="${SSH_PORT:-$(detect_ssh_port)}"
SSH_PORT="$(printf '%s' "$SSH_PORT" | tr -d '[:space:]')"
case "$SSH_PORT" in
  ''|*[!0-9,]*)
    echo "Invalid SSH port: '$SSH_PORT' (expected digits, e.g. 22 or 22,2222)." >&2
    exit 1
    ;;
esac

BANACTION="$(detect_banaction)"
BANACTION_ALLPORTS="$(detect_banaction_allports "$BANACTION")"
CLIENT_IP="$(detect_client_ip | tr -d '[:space:]')"
IGNORE_IP="$(printf '%s' "$IGNORE_IP" | tr -d '[:space:]')"

# Values are interpolated into the jail config, so reject anything that is not
# a plain address/CIDR list before it can add or override a directive.
for entry in ${CLIENT_IP:-} ${IGNORE_IP:-}; do
  case "$entry" in
    *[!0-9a-fA-F:.,/]*|*/*/*)
      echo "Invalid IP entry: '$entry' (expected address or CIDR)." >&2
      exit 1
      ;;
  esac
done

IGNORE_LIST="127.0.0.1/8 ::1"
if [ -n "$CLIENT_IP" ]; then
  IGNORE_LIST="$IGNORE_LIST $CLIENT_IP"
fi
if [ -n "$IGNORE_IP" ]; then
  IGNORE_LIST="$IGNORE_LIST $IGNORE_IP"
fi

echo "    ssh port  : $SSH_PORT"
echo "    sshd mode : $SSHD_MODE"
echo "    banaction : $BANACTION (all-ports: $BANACTION_ALLPORTS)"
echo "    logtarget : $FAIL2BAN_LOG"
echo "    ignoreip  : $IGNORE_LIST"

if [ -z "$CLIENT_IP" ]; then
  echo "    WARNING: this session's client address could not be detected, so it is" >&2
  echo "             NOT whitelisted. With $MAXRETRY attempts in $FINDTIME -> $BANTIME" >&2
  echo "             ($RECIDIVE_MAXRETRY bans -> $RECIDIVE_BANTIME) you can lock yourself out." >&2
  echo "             Re-run with IGNORE_IP=<your address> if you administer over SSH." >&2
fi

echo "==> Ensuring sshd logs failed attempts"
if [ "$(sshd_effective loglevel | head -n1)" = "QUIET" ]; then
  install -d -m 0755 /etc/ssh/sshd_config.d
  printf 'LogLevel VERBOSE\n' > /etc/ssh/sshd_config.d/99-fail2ban-loglevel.conf
  if sshd -t 2>/dev/null; then
    systemctl reload ssh 2>/dev/null || systemctl reload sshd 2>/dev/null || \
      echo "    warning: could not reload sshd; the LogLevel change applies after its next restart" >&2
    echo "    set LogLevel VERBOSE"
  else
    rm -f /etc/ssh/sshd_config.d/99-fail2ban-loglevel.conf
    echo "    skipped: config test failed" >&2
  fi
else
  echo "    already sufficient"
fi

echo "==> Writing $JAIL_LOCAL"
if [ -f "$JAIL_LOCAL" ]; then
  cp "$JAIL_LOCAL" "${JAIL_LOCAL}.bak.$(date +%s)"
fi

# jail.d/*.conf is read BEFORE jail.local, so our settings win and entries that
# other packages or the admin put there (nginx, postfix, ...) keep working.
# They are therefore left untouched - only reported.
if [ -d "$JAIL_D" ]; then
  existing_jails="$(find "$JAIL_D" -maxdepth 1 -type f -name '*.conf' -printf '%f ' 2>/dev/null || true)"
  if [ -n "${existing_jails:-}" ]; then
    echo "    keeping existing entries in $JAIL_D: $existing_jails"
  fi
fi

cat > "$JAIL_LOCAL" <<EOF
[DEFAULT]
ignoreip = $IGNORE_LIST
bantime  = $BANTIME
findtime = $FINDTIME
maxretry = $MAXRETRY
backend  = systemd
banaction = $BANACTION
banaction_allports = $BANACTION_ALLPORTS
# Repeat offenders are handled by the [recidive] jail below. These options only
# take effect if bantime.increment is switched to true (maxtime caps the growth).
bantime.increment = false
bantime.maxtime = 5w

[sshd]
enabled  = true
port     = $SSH_PORT
mode     = $SSHD_MODE
# Inert with the systemd backend (fail2ban uses the filter's journalmatch and
# ignores logpath); kept for hosts that switch this jail to a file backend.
logpath  = %(sshd_log)s

# fail2ban writes its ban notices to $FAIL2BAN_LOG (logtarget), not to the
# journal, so this jail must read that file: with the systemd backend inherited
# from [DEFAULT] it would monitor a journal that never contains a "Ban" line and
# would silently never ban anyone.
[recidive]
enabled  = true
backend  = polling
journalmatch =
logpath  = $FAIL2BAN_LOG
findtime = $RECIDIVE_FINDTIME
maxretry = $RECIDIVE_MAXRETRY
bantime  = $RECIDIVE_BANTIME
EOF

# A file backend refuses to start fail2ban when its logpath does not exist yet.
install -d -m 0755 "$(dirname "$FAIL2BAN_LOG")"
touch "$FAIL2BAN_LOG"

LOG_SIZE_BEFORE="$(wc -c < "$FAIL2BAN_LOG" 2>/dev/null | tr -d '[:space:]')"
LOG_SIZE_BEFORE="${LOG_SIZE_BEFORE:-0}"

echo "==> Enabling and restarting fail2ban"
systemctl enable fail2ban >/dev/null 2>&1 || true
if ! systemctl restart fail2ban; then
  echo "fail2ban failed to start. Recent log:" >&2
  journalctl -u fail2ban -n 30 --no-pager >&2 || true
  exit 1
fi

echo "==> Verifying"
# fail2ban can be "active" while a jail silently failed to load, so wait for the
# jails themselves and fail loudly instead of printing a success summary.
jails_ok=false
tries=0
while [ "$tries" -lt 10 ]; do
  if systemctl is-active --quiet fail2ban \
     && fail2ban-client status sshd >/dev/null 2>&1 \
     && fail2ban-client status recidive >/dev/null 2>&1; then
    jails_ok=true
    break
  fi
  tries=$((tries + 1))
  sleep 1
done

if [ "$jails_ok" != true ]; then
  echo "fail2ban is not enforcing: the sshd/recidive jails did not come up." >&2
  journalctl -u fail2ban -n 30 --no-pager >&2 || true
  fail2ban-client status >&2 || true
  exit 1
fi

# An active jail proves nothing about [recidive]: it only bans when the log it
# polls receives lines. Wait for the restart to actually write there.
log_written=false
tries=0
while [ "$tries" -lt 5 ]; do
  LOG_SIZE_NOW="$(wc -c < "$FAIL2BAN_LOG" 2>/dev/null | tr -d '[:space:]')"
  if [ "${LOG_SIZE_NOW:-0}" -gt "$LOG_SIZE_BEFORE" ]; then
    log_written=true
    break
  fi
  tries=$((tries + 1))
  sleep 1
done

if [ "$log_written" != true ]; then
  echo "fail2ban is not writing to $FAIL2BAN_LOG, so the [recidive] jail would" >&2
  echo "never observe a ban. Check 'logtarget' in fail2ban.conf/fail2ban.local." >&2
  journalctl -u fail2ban -n 30 --no-pager >&2 || true
  exit 1
fi

fail2ban-client status
fail2ban-client status sshd
fail2ban-client status recidive

echo
echo "fail2ban configured with a strict profile."
echo "Jails:            sshd (port $SSH_PORT, mode $SSHD_MODE), recidive"
echo "Ban policy:       $MAXRETRY failures in $FINDTIME -> $BANTIME"
echo "Repeat offenders: $RECIDIVE_MAXRETRY bans in $RECIDIVE_FINDTIME -> $RECIDIVE_BANTIME"
echo "Recidive source:  $FAIL2BAN_LOG"
echo "Never banned:     $IGNORE_LIST"
echo
echo "Check:   fail2ban-client status sshd"
echo "Unban:   fail2ban-client set sshd unbanip <IP>"
echo "Config:  $JAIL_LOCAL"