#!/usr/bin/env bash
# Configure Ubuntu UFW for the isolated MTProxyMax deployment.
# - Allows all proxy TCP ports in both directions.
# - Detects and allows the active/configured SSH TCP port(s) first.
# - Applies rules to IPv4 and IPv6 through UFW.
# - Blocks TCP and UDP port 111 in both directions with highest priority.

set -Eeuo pipefail

readonly BLOCKED_PORT=111
readonly -a PROXY_PORTS=(21 80 443 8080 2083 2086 2096 2082 2095 586)

log() {
    printf '[ufw-mtproxy] %s\n' "$*"
}

die() {
    printf '[ufw-mtproxy] ERROR: %s\n' "$*" >&2
    exit 1
}

on_error() {
    local exit_code=$?
    printf '[ufw-mtproxy] ERROR: failed at line %s (exit %s).\n' "${BASH_LINENO[0]:-unknown}" "$exit_code" >&2
    exit "$exit_code"
}
trap on_error ERR

[ "$(id -u)" -eq 0 ] || die "Run this script as root (for example: sudo bash $0)."
[ -r /etc/os-release ] || die "Cannot identify the operating system."

# shellcheck disable=SC1091
. /etc/os-release
case "${ID:-}:${ID_LIKE:-}" in
    ubuntu:*|*:ubuntu*) ;;
    *) die "This script supports Ubuntu only; detected ${PRETTY_NAME:-unknown}." ;;
esac

if ! command -v ufw >/dev/null 2>&1; then
    log "UFW is not installed; installing it with apt..."
    export DEBIAN_FRONTEND=noninteractive
    apt-get update
    apt-get install -y ufw
else
    log "UFW is already installed."
fi

command -v ufw >/dev/null 2>&1 || die "UFW installation did not complete successfully."

# UFW generates matching IPv6 rules only when IPV6=yes. Preserve one backup
# of the original defaults file before changing this setting.
readonly UFW_DEFAULTS=/etc/default/ufw
[ -f "$UFW_DEFAULTS" ] || die "Missing ${UFW_DEFAULTS}."
if [ ! -f "${UFW_DEFAULTS}.before-mtproxy" ]; then
    cp -a "$UFW_DEFAULTS" "${UFW_DEFAULTS}.before-mtproxy"
fi
if grep -Eq '^[[:space:]]*IPV6=' "$UFW_DEFAULTS"; then
    sed -Ei 's/^[[:space:]]*IPV6=.*/IPV6=yes/' "$UFW_DEFAULTS"
else
    printf '\nIPV6=yes\n' >> "$UFW_DEFAULTS"
fi
log "UFW IPv6 rule generation is enabled."

declare -A SSH_PORT_SET=()

add_ssh_port() {
    local port="${1:-}"
    if [[ "$port" =~ ^[0-9]+$ ]] && [ "$port" -ge 1 ] && [ "$port" -le 65535 ]; then
        SSH_PORT_SET["$port"]=1
    fi
}

# The fourth field is the local server port of the current SSH connection.
if [ -n "${SSH_CONNECTION:-}" ]; then
    read -r _ssh_client_ip _ssh_client_port _ssh_server_ip _ssh_server_port <<< "$SSH_CONNECTION"
    add_ssh_port "${_ssh_server_port:-}"
fi

# Ask sshd for its effective configured ports, including sshd_config.d files.
SSHD_BIN=""
if command -v sshd >/dev/null 2>&1; then
    SSHD_BIN=$(command -v sshd)
elif [ -x /usr/sbin/sshd ]; then
    SSHD_BIN=/usr/sbin/sshd
fi
if [ -n "$SSHD_BIN" ]; then
    while read -r ssh_port; do
        add_ssh_port "$ssh_port"
    done < <("$SSHD_BIN" -T 2>/dev/null | awk '$1 == "port" { print $2 }' || true)
fi

# Also inspect live sshd listeners. This protects custom ports even if the
# current session/environment does not expose SSH_CONNECTION.
if command -v ss >/dev/null 2>&1; then
    while read -r ssh_port; do
        add_ssh_port "$ssh_port"
    done < <(ss -H -ltnp 2>/dev/null | awk '
        /sshd/ {
            endpoint=$4
            sub(/^.*:/, "", endpoint)
            if (endpoint ~ /^[0-9]+$/) print endpoint
        }
    ' || true)
fi

# Port 22 is the safe fallback when sshd is not installed yet or cannot expose
# its effective configuration. On an active custom-port SSH session, the
# session port above is used instead.
if [ ${#SSH_PORT_SET[@]} -eq 0 ]; then
    add_ssh_port 22
fi

mapfile -t SSH_PORTS < <(printf '%s\n' "${!SSH_PORT_SET[@]}" | sort -n)
log "SSH TCP port(s) that will remain allowed: ${SSH_PORTS[*]}"

if [ -n "${SSH_PORT_SET[$BLOCKED_PORT]:-}" ]; then
    die "SSH is using port ${BLOCKED_PORT}; refusing to block it because that would disconnect administrative access. Move SSH to another port first."
fi

declare -A ALLOW_PORT_SET=()
for port in "${PROXY_PORTS[@]}" "${SSH_PORTS[@]}"; do
    ALLOW_PORT_SET["$port"]=1
done
mapfile -t ALLOW_PORTS < <(printf '%s\n' "${!ALLOW_PORT_SET[@]}" | sort -n)

# Add SSH/proxy rules before enabling UFW so the current administration session
# cannot be locked out. MTProxy uses TCP; UDP is intentionally not exposed.
for port in "${ALLOW_PORTS[@]}"; do
    if [[ " ${SSH_PORTS[*]} " == *" ${port} "* ]]; then
        rule_comment="SSH TCP ${port}"
    else
        rule_comment="MTProxy TCP ${port}"
    fi
    ufw allow in proto tcp to any port "$port" comment "$rule_comment"
    ufw allow out proto tcp to any port "$port" comment "$rule_comment"
done

# Port 111 is commonly used by rpcbind. Insert the denies at position 1 so
# they take precedence over an older broad allow rule. Block both TCP and UDP,
# and block packets whose source or destination is port 111.
for protocol in tcp udp; do
    ufw insert 1 deny in proto "$protocol" to any port "$BLOCKED_PORT" comment "Block rpcbind destination ${protocol}"
    ufw insert 1 deny in proto "$protocol" from any port "$BLOCKED_PORT" comment "Block rpcbind source ${protocol}"
    ufw insert 1 deny out proto "$protocol" to any port "$BLOCKED_PORT" comment "Block rpcbind destination ${protocol}"
    ufw insert 1 deny out proto "$protocol" from any port "$BLOCKED_PORT" comment "Block rpcbind source ${protocol}"
done

# Enabling an already active UFW installation is safe and keeps this script
# idempotent. Reload ensures an IPV6=no -> yes change takes effect immediately.
ufw --force enable
ufw reload

log "Firewall configuration completed. Current rules:"
ufw status verbose
echo
ufw status numbered

if command -v sysctl >/dev/null 2>&1 && [ "$(sysctl -n net.ipv6.conf.all.disable_ipv6 2>/dev/null || echo 0)" = "1" ]; then
    log "WARNING: IPv6 is disabled at kernel level. UFW IPv6 rules were created, but IPv6 traffic cannot be active until the kernel setting is enabled."
fi

log "Allowed TCP ports (input/output): ${ALLOW_PORTS[*]}"
log "Blocked completely (input/output, TCP/UDP): ${BLOCKED_PORT}"