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

# Ubuntu Firewall Rules Validation Script
#
# Purpose:
#   Validate host firewall state and common exposure paths when using UFW and/or nftables.
#   This script is non-destructive: it only inspects configuration and runs connectivity checks
#   you explicitly request via arguments.
#
# Usage:
#   ./ubuntu-firewall-validate.sh [options]
#
# Options:
#   --check-ufw                 Check whether UFW is installed, active, and readable.
#   --check-nftables            Check whether nftables is installed and list active ruleset.
#   --check-port HOST:PORT      Test TCP connectivity to HOST:PORT from this machine.
#   --check-ssh HOST[:PORT]     Convenience shortcut for TCP connectivity to SSH port 22 if PORT omitted.
#   --expect-open HOST:PORT     Mark the TCP check as expected to succeed.
#   --expect-closed HOST:PORT   Mark the TCP check as expected to fail.
#   --help                      Show help.
#
# Examples:
#   ./ubuntu-firewall-validate.sh --check-ufw --check-nftables
#   ./ubuntu-firewall-validate.sh --check-ssh 192.0.2.10:22 --expect-open 192.0.2.10:22
#   ./ubuntu-firewall-validate.sh --check-port 198.51.100.20:443 --expect-open 198.51.100.20:443

print_help() {
  cat <<'EOF'
Ubuntu Firewall Rules Validation Script

Validate UFW/nftables state and optional TCP reachability checks.

Options:
  --check-ufw                 Check UFW status
  --check-nftables            Check nftables ruleset
  --check-port HOST:PORT      Test TCP connectivity to HOST:PORT
  --check-ssh HOST[:PORT]     Test SSH port; default port is 22
  --expect-open HOST:PORT     Expect TCP check to succeed
  --expect-closed HOST:PORT   Expect TCP check to fail
  --help                      Show this help
EOF
}

have_cmd() {
  command -v "$1" >/dev/null 2>&1
}

check_ufw() {
  echo "== UFW =="
  if ! have_cmd ufw; then
    echo "UFW: not installed"
    return 0
  fi

  if ufw status >/tmp/ufw-status.$$ 2>/tmp/ufw-status.err; then
    cat /tmp/ufw-status.$$
  else
    echo "UFW: installed, but status check failed"
    sed 's/^/  /' /tmp/ufw-status.err || true
  fi

  rm -f /tmp/ufw-status.$$ /tmp/ufw-status.err
}

check_nftables() {
  echo "== nftables =="
  if ! have_cmd nft; then
    echo "nftables: not installed"
    return 0
  fi

  if nft list ruleset; then
    :
  else
    echo "nftables: ruleset listing failed"
  fi
}

check_tcp() {
  local host_port="$1"
  local expected="$2"
  local host port

  if [[ "$host_port" != *:* ]]; then
    echo "Invalid HOST:PORT value: $host_port" >&2
    return 2
  fi

  host="${host_port%:*}"
  port="${host_port##*:}"

  if [[ -z "$host" || -z "$port" ]]; then
    echo "Invalid HOST:PORT value: $host_port" >&2
    return 2
  fi

  echo "== TCP check: $host:$port =="

  local ok=0
  if have_cmd nc; then
    if nc -z -w 3 "$host" "$port" >/dev/null 2>&1; then
      ok=1
    fi
  elif have_cmd timeout && have_cmd bash; then
    if timeout 3 bash -c ">/dev/tcp/$host/$port" >/dev/null 2>&1; then
      ok=1
    fi
  else
    echo "No suitable TCP test tool found (install nc/netcat or use bash with /dev/tcp)." >&2
    return 2
  fi

  if [[ "$expected" == "open" ]]; then
    if [[ "$ok" -eq 1 ]]; then
      echo "Result: reachable as expected"
    else
      echo "Result: unreachable, but expected open" >&2
      return 1
    fi
  elif [[ "$expected" == "closed" ]]; then
    if [[ "$ok" -eq 1 ]]; then
      echo "Result: reachable, but expected closed" >&2
      return 1
    else
      echo "Result: unreachable as expected"
    fi
  else
    if [[ "$ok" -eq 1 ]]; then
      echo "Result: reachable"
    else
      echo "Result: unreachable"
    fi
  fi
}

main() {
  local -a port_checks=()
  local -a expected_open=()
  local -a expected_closed=()
  local check_ufw_flag=0
  local check_nft_flag=0

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --check-ufw)
        check_ufw_flag=1
        shift
        ;;
      --check-nftables)
        check_nft_flag=1
        shift
        ;;
      --check-port)
        [[ $# -ge 2 ]] || { echo "--check-port requires HOST:PORT" >&2; exit 2; }
        port_checks+=("$2")
        shift 2
        ;;
      --check-ssh)
        [[ $# -ge 2 ]] || { echo "--check-ssh requires HOST[:PORT]" >&2; exit 2; }
        if [[ "$2" == *:* ]]; then
          port_checks+=("$2")
        else
          port_checks+=("$2:22")
        fi
        shift 2
        ;;
      --expect-open)
        [[ $# -ge 2 ]] || { echo "--expect-open requires HOST:PORT" >&2; exit 2; }
        expected_open+=("$2")
        shift 2
        ;;
      --expect-closed)
        [[ $# -ge 2 ]] || { echo "--expect-closed requires HOST:PORT" >&2; exit 2; }
        expected_closed+=("$2")
        shift 2
        ;;
      --help|-h)
        print_help
        exit 0
        ;;
      *)
        echo "Unknown argument: $1" >&2
        print_help >&2
        exit 2
        ;;
    esac
  done

  local did_anything=0

  if [[ "$check_ufw_flag" -eq 1 ]]; then
    did_anything=1
    check_ufw
    echo
  fi

  if [[ "$check_nft_flag" -eq 1 ]]; then
    did_anything=1
    check_nftables
    echo
  fi

  local item
  for item in "${expected_open[@]}"; do
    did_anything=1
    check_tcp "$item" "open"
    echo
  done

  for item in "${expected_closed[@]}"; do
    did_anything=1
    check_tcp "$item" "closed"
    echo
  done

  for item in "${port_checks[@]}"; do
    did_anything=1
    check_tcp "$item" "unknown"
    echo
  done

  if [[ "$did_anything" -eq 0 ]]; then
    print_help
    exit 0
  fi
}

main "$@"