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

# UFW Safe Setup and Validation Script
#
# Purpose:
#   Help configure UFW on Ubuntu Server in a cautious, repeatable way.
#   This script avoids destructive defaults, never removes rules automatically,
#   and prompts before enabling the firewall or applying changes.
#
# What it does:
#   - Checks whether ufw and sudo are available
#   - Shows current UFW status
#   - Optionally installs ufw if missing
#   - Optionally enables UFW after confirming management access is allowed
#   - Applies a default deny incoming / allow outgoing posture
#   - Adds explicit allow rules from arguments
#   - Shows numbered rules for review
#
# Examples:
#   # Review current status only
#   ./ufw-safe-setup.sh
#
#   # Install ufw if needed, allow SSH and HTTPS, then enable UFW
#   ./ufw-safe-setup.sh --install --enable --allow OpenSSH --allow 443/tcp
#
#   # Allow SSH only from an admin subnet and HTTPS from anywhere
#   ./ufw-safe-setup.sh \
#     --allow "from 203.0.113.0/24 to any port 22 proto tcp" \
#     --allow 443/tcp --enable
#
# Notes:
#   - Run from a console or with a known recovery path.
#   - Adjust rules to match your environment.
#   - This script does not delete rules.

usage() {
  cat <<'EOF'
Usage: ufw-safe-setup.sh [options]

Options:
  --install                 Install ufw if it is not present
  --enable                  Prompt before enabling ufw
  --allow RULE              Add an allow rule (repeatable)
  --default-deny-incoming   Set default incoming policy to deny
  --default-allow-outgoing  Set default outgoing policy to allow
  --show                    Show UFW status and numbered rules (default action)
  -h, --help                Show this help message

Examples:
  ufw-safe-setup.sh --install --default-deny-incoming --default-allow-outgoing --allow OpenSSH --allow 443/tcp --enable
  ufw-safe-setup.sh --allow "from 198.51.100.25 to any port 22 proto tcp" --enable
EOF
}

require_root_or_sudo() {
  if [[ ${EUID:-$(id -u)} -ne 0 ]]; then
    if ! command -v sudo >/dev/null 2>&1; then
      echo "Error: root privileges or sudo are required." >&2
      exit 1
    fi
  fi
}

run() {
  if [[ ${EUID:-$(id -u)} -eq 0 ]]; then
    "$@"
  else
    sudo "$@"
  fi
}

confirm() {
  local prompt="$1"
  read -r -p "$prompt [y/N]: " reply
  [[ "$reply" =~ ^[Yy]$ ]]
}

show_status() {
  echo
  echo "== UFW status =="
  if command -v ufw >/dev/null 2>&1; then
    run ufw status verbose || true
    echo
    echo "== UFW numbered rules =="
    run ufw status numbered || true
  else
    echo "ufw is not installed."
  fi
}

INSTALL=false
ENABLE=false
SET_DENY_IN=false
SET_ALLOW_OUT=false
SHOW_ONLY=false
declare -a ALLOW_RULES=()

while [[ $# -gt 0 ]]; do
  case "$1" in
    --install)
      INSTALL=true
      ;;
    --enable)
      ENABLE=true
      ;;
    --allow)
      shift
      [[ $# -gt 0 ]] || { echo "Error: --allow requires a rule argument." >&2; exit 1; }
      ALLOW_RULES+=("$1")
      ;;
    --default-deny-incoming)
      SET_DENY_IN=true
      ;;
    --default-allow-outgoing)
      SET_ALLOW_OUT=true
      ;;
    --show)
      SHOW_ONLY=true
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Error: unknown option: $1" >&2
      usage
      exit 1
      ;;
  esac
  shift
 done

require_root_or_sudo

if ! command -v ufw >/dev/null 2>&1; then
  if [[ "$INSTALL" == true ]]; then
    echo "Installing ufw..."
    run apt update
    run apt install -y ufw
  else
    echo "ufw is not installed. Re-run with --install if you want this script to install it."
    exit 0
  fi
fi

show_status

if [[ "$SHOW_ONLY" == true && "$ENABLE" == false && "$SET_DENY_IN" == false && "$SET_ALLOW_OUT" == false && ${#ALLOW_RULES[@]} -eq 0 ]]; then
  exit 0
fi

if [[ "$SET_DENY_IN" == true ]]; then
  echo "Setting default incoming policy to deny..."
  run ufw default deny incoming
fi

if [[ "$SET_ALLOW_OUT" == true ]]; then
  echo "Setting default outgoing policy to allow..."
  run ufw default allow outgoing
fi

for rule in "${ALLOW_RULES[@]}"; do
  echo "Adding allow rule: $rule"
  run ufw allow "$rule"
done

show_status

if [[ "$ENABLE" == true ]]; then
  echo
  echo "Before enabling UFW, confirm that your management access is allowed."
  echo "Review the rules above carefully."
  if confirm "Enable UFW now?"; then
    run ufw enable
    echo "UFW enabled."
  else
    echo "Skipped enabling UFW."
  fi
fi

echo
echo "Final rule review:"
run ufw status numbered || true

echo
echo "Done. Verify access from a remote client before production use."