#!/bin/bash
# dfw - open and close ports in the Das OS server firewall (nftables).
#
# /etc/nftables.conf defines the "das" table with the sets tcp_ports and
# udp_ports; dfw keeps their elements in /etc/nftables.d/dfw.nft, which that
# file includes, and reloads only the "das" table. SSH is always open.
set -euo pipefail

# shellcheck source=helpers.sh
. "${DAS_HELPERS:-/usr/lib/das/helpers.sh}"

conf=/etc/nftables.conf
ports_file=/etc/nftables.d/dfw.nft

usage() {
  if [[ -n $uk ]]; then
    cat <<'EOF'
Використання: dfw КОМАНДА [ПОРТ...]

  status                  стан брандмауера й відкриті порти (типова команда)
  allow ПОРТ...           відкрити порти для вхідних з'єднань
  deny ПОРТ...            закрити порти, відкриті через allow
  reload                  перечитати /etc/nftables.conf після редагування

ПОРТ: 80, 80/tcp, 53/udp, 8000-8010/tcp або назва служби (http, https).
Без /udp мається на увазі tcp. SSH відкритий завжди.

Порти, які публікують контейнери Docker (-p 8080:80), відкриває сам Docker,
dfw на них не впливає.
EOF
  else
    cat <<'EOF'
Usage: dfw COMMAND [PORT...]

  status                  firewall state and open ports (the default)
  allow PORT...           open ports for incoming connections
  deny PORT...            close ports opened with allow
  reload                  load /etc/nftables.conf again after editing it

PORT: 80, 80/tcp, 53/udp, 8000-8010/tcp or a service name (http, https).
Without /udp it means tcp. SSH is always open.

Ports that Docker containers publish (-p 8080:80) are opened by Docker
itself, dfw does not affect them.
EOF
  fi
}

check_config() {
  grep -qs 'set tcp_ports' "$conf" ||
    die "$conf не з Das OS Server (немає набору tcp_ports); dfw працює лише з ним" \
      "$conf is not the Das OS Server one (no tcp_ports set); dfw only works with that"
}

# Opened ports of one protocol from the ports file, one per line
saved_ports() {
  [[ -r $ports_file ]] || return 0
  sed -n "s/^add element inet das $1_ports { \(.*\) }$/\1/p" "$ports_file" | tr -d ' ' | tr ',' '\n'
}

# parse_port SPEC: prints "PORT PROTO" with PORT a number or a FIRST-LAST range
parse_port() {
  local spec=$1 port proto=tcp first last
  if [[ $spec == */* ]]; then
    port=${spec%/*}
    proto=${spec##*/}
  else
    port=$spec
  fi
  [[ $proto == tcp || $proto == udp ]] || die "'$spec': протокол tcp або udp" "'$spec': the protocol is tcp or udp"
  if [[ $port =~ ^[a-z][a-z0-9-]*$ ]]; then
    # A service name from /etc/services
    first=$(getent services "$port/$proto" | awk '{ print $2 }')
    [[ -n $first ]] || die "'$spec': невідома служба $port/$proto" "'$spec': unknown service $port/$proto"
    port=${first%/*}
  fi
  if [[ $port =~ ^([0-9]+)-([0-9]+)$ ]]; then
    first=${BASH_REMATCH[1]}
    last=${BASH_REMATCH[2]}
  elif [[ $port =~ ^[0-9]+$ ]]; then
    first=$port
    last=$port
  else
    die "'$spec': порт — число, діапазон 8000-8010 або назва служби" \
      "'$spec': a port is a number, a range like 8000-8010 or a service name"
  fi
  ((first >= 1 && last <= 65535 && first <= last)) ||
    die "'$spec': порти від 1 до 65535" "'$spec': ports are 1 to 65535"
  if ((first == last)); then
    echo "$first $proto"
  else
    echo "$first-$last $proto"
  fi
}

firewall_on() {
  [[ -e /etc/runit/runsvdir/default/nftables ]]
}

# save TCP_PORTS UDP_PORTS (newline-separated): writes the ports file, checks
# the whole configuration and loads it, restoring the old file on an error
save() {
  local tcp udp backup=""
  # sort -n -u would take 80 and 80-90 for duplicates
  tcp=$(sort -u <<<"$1" | sort -n | grep . | paste -sd, - | sed 's/,/, /g' || true)
  udp=$(sort -u <<<"$2" | sort -n | grep . | paste -sd, - | sed 's/,/, /g' || true)
  mkdir -p "${ports_file%/*}"
  if [[ -e $ports_file ]]; then
    backup=$(mktemp)
    cp "$ports_file" "$backup"
  fi
  {
    echo "# Ports opened with dfw; change them with dfw allow/deny"
    [[ -z $tcp ]] || echo "add element inet das tcp_ports { $tcp }"
    [[ -z $udp ]] || echo "add element inet das udp_ports { $udp }"
  } >"$ports_file"
  if ! nft -c -f "$conf"; then
    if [[ -n $backup ]]; then mv "$backup" "$ports_file"; else rm -f "$ports_file"; fi
    die "nftables не прийняв правила, нічого не змінено" "nftables rejected the rules, nothing changed"
  fi
  rm -f "$backup"
  if firewall_on; then
    nft -f "$conf"
  else
    warn "брандмауер вимкнений, зміни збережено; увімкнути: dsv enable nftables" \
      "the firewall is off, changes are saved; to turn it on: dsv enable nftables"
  fi
}

cmd_allow() {
  local spec port proto tcp udp
  tcp=$(saved_ports tcp)
  udp=$(saved_ports udp)
  for spec in "$@"; do
    read -r port proto < <(parse_port "$spec")
    if [[ $port == 22 && $proto == tcp ]]; then
      say "22/tcp (SSH) відкритий завжди" "22/tcp (SSH) is always open"
      continue
    fi
    if grep -qx "$port" <<<"$(if [[ $proto == tcp ]]; then echo "$tcp"; else echo "$udp"; fi)"; then
      say "$port/$proto: уже відкритий" "$port/$proto: already open"
      continue
    fi
    if [[ $proto == tcp ]]; then tcp+=$'\n'$port; else udp+=$'\n'$port; fi
    say "$port/$proto: відкрито" "$port/$proto: open"
  done
  save "$tcp" "$udp"
}

cmd_deny() {
  local spec port proto tcp udp list
  tcp=$(saved_ports tcp)
  udp=$(saved_ports udp)
  for spec in "$@"; do
    read -r port proto < <(parse_port "$spec")
    if [[ $port == 22 && $proto == tcp ]]; then
      die "SSH не закривається через dfw, щоб не втратити доступ до сервера; правило — у $conf" \
        "SSH is not closed with dfw, so as not to lose access to the server; the rule is in $conf"
    fi
    if [[ $proto == tcp ]]; then list=$tcp; else list=$udp; fi
    if ! grep -qx "$port" <<<"$list"; then
      say "$port/$proto: не був відкритий через dfw" "$port/$proto: was not opened with dfw"
      continue
    fi
    list=$(grep -vx "$port" <<<"$list" || true)
    if [[ $proto == tcp ]]; then tcp=$list; else udp=$list; fi
    say "$port/$proto: закрито" "$port/$proto: closed"
  done
  save "$tcp" "$udp"
}

cmd_status() {
  local p state banned ports
  if firewall_on && nft list table inet das >/dev/null 2>&1; then
    state=$(say "увімкнений" "on")
  else
    state=$(say "ВИМКНЕНИЙ (dsv enable nftables)" "OFF (dsv enable nftables)")
  fi
  say "Брандмауер: $state" "Firewall: $state"
  say "Вхідні з'єднання закриті, крім:" "Incoming connections are closed except:"
  echo "  22/tcp (ssh)"
  for p in tcp udp; do
    while read -r ports; do
      [[ -n $ports ]] && echo "  $ports/$p"
    done < <(saved_ports "$p")
  done

  # Docker opens its published ports itself
  if command -v docker >/dev/null && ports=$(docker ps --format '{{.Names}}: {{.Ports}}' 2>/dev/null) &&
    grep -q -- '->' <<<"$ports"; then
    say "Опубліковані контейнерами Docker (відкриті незалежно від dfw):" \
      "Published by Docker containers (open regardless of dfw):"
    grep -- '->' <<<"$ports" | sed 's/^/  /'
  fi

  if banned=$(fail2ban-client status sshd 2>/dev/null | sed -n 's/.*Currently banned:[[:space:]]*//p') &&
    [[ -n $banned && $banned != 0 ]]; then
    say "fail2ban: заблоковано адрес за невдалі входи SSH: $banned" \
      "fail2ban: addresses banned for failed SSH logins: $banned"
  fi
}

cmd=${1:-status}
[[ $# -gt 0 ]] && shift

case $cmd in
  status | list | ls)
    need_root status
    check_config
    cmd_status
    ;;
  allow | open)
    [[ $# -gt 0 ]] || die "вкажіть порт: dfw allow 80/tcp" "give a port: dfw allow 80/tcp"
    need_root allow "$@"
    check_config
    cmd_allow "$@"
    ;;
  deny | close)
    [[ $# -gt 0 ]] || die "вкажіть порт: dfw deny 80/tcp" "give a port: dfw deny 80/tcp"
    need_root deny "$@"
    check_config
    cmd_deny "$@"
    ;;
  reload)
    need_root reload
    nft -c -f "$conf"
    if firewall_on; then
      nft -f "$conf"
      say "правила завантажено" "rules loaded"
    else
      say "правила правильні, але брандмауер вимкнений: dsv enable nftables" \
        "the rules are valid, but the firewall is off: dsv enable nftables"
    fi
    ;;
  -h | --help | help) usage ;;
  *)
    usage >&2
    exit 1
    ;;
esac