#!/usr/bin/env bash

sshConfigDropIn="/etc/ssh/sshd_config.d/99-stormux-server.conf"

ufw_installed() {
    pacman -Q ufw &> /dev/null
}

ufw_status_output() {
    # `sudoFlags` is initialized by the main launcher before sourcing this file.
    # shellcheck disable=SC2154
    sudo "${sudoFlags[@]}" ufw status 2>&1
}

sudo_ready_noninteractive() {
    # `sudoFlags` is initialized by the main launcher before sourcing this file.
    # shellcheck disable=SC2154
    sudo "${sudoFlags[@]}" -n true &> /dev/null
}

ufw_status_output_noninteractive() {
    # `sudoFlags` is initialized by the main launcher before sourcing this file.
    # shellcheck disable=SC2154
    sudo "${sudoFlags[@]}" -n ufw status 2>&1
}

warn_sudo_authentication() {
    if sudo_ready_noninteractive; then
        return 0
    fi

    msgbox "This firewall action may require sudo authentication. If you hear the password prompt sound, enter your password and press Enter."
    return 0
}

firewall_reboot_required() {
    [[ ! -d "/lib/modules/$(uname -r)" ]]
}

ensure_firewall_backend_ready() {
    if firewall_reboot_required; then
        msgbox "The running kernel does not match the installed modules. Reboot this server before managing the firewall."
        return 1
    fi

    return 0
}

firewall_enabled() {
    local statusText=""

    if ! ufw_installed; then
        return 1
    fi

    if ! ensure_firewall_backend_ready; then
        return 1
    fi

    if ! sudo_ready_noninteractive; then
        return 1
    fi

    statusText="$(ufw_status_output_noninteractive)"
    grep -q '^Status: active$' <<< "$statusText"
}

ensure_ufw() {
    if ufw_installed; then
        return 0
    fi

    if [[ "$(yesno "ufw is not installed. Install it now and continue?")" != "Yes" ]]; then
        msgbox "Firewall action cancelled."
        return 1
    fi

    warn_sudo_authentication
    if ! install_package ufw; then
        msgbox "Failed to install ufw."
        return 1
    fi

    msgbox "ufw installed. The firewall is not active until you choose Enable firewall."
    return 0
}

valid_port() {
    local portValue="$1"
    [[ "$portValue" =~ ^[0-9]+$ ]] && (( portValue >= 1 && portValue <= 65535 ))
}

valid_port_rule() {
    local ruleValue="$1"

    if [[ "$ruleValue" =~ ^[0-9]+$ ]]; then
        valid_port "$ruleValue"
        return $?
    fi

    if [[ "$ruleValue" =~ ^([0-9]+)/(tcp|udp)$ ]]; then
        valid_port "${BASH_REMATCH[1]}"
        return $?
    fi

    return 1
}

resolve_ssh_port() {
    local configuredPort=""

    if [[ -r "$sshConfigDropIn" ]]; then
        configuredPort="$(awk '
            BEGIN { IGNORECASE = 1 }
            /^[[:space:]]*Port[[:space:]]+[0-9]+([[:space:]]*#.*)?$/ {
                port = $2
            }
            END {
                if (port != "") {
                    print port
                }
            }
        ' "$sshConfigDropIn")"
    fi

    if valid_port "$configuredPort"; then
        printf '%s\n' "$configuredPort"
        return 0
    fi

    configuredPort="$(inputbox "Unable to confirm the SSH port automatically. Enter the SSH port to allow before changing the firewall.")" || return 1
    if valid_port "$configuredPort"; then
        printf '%s\n' "$configuredPort"
        return 0
    fi

    msgbox "A valid SSH port is required before enabling firewall rules."
    return 1
}

allow_rule() {
    local ruleValue="$1"
    local description="$2"

    # `sudoFlags` is initialized by the main launcher before sourcing this file.
    # shellcheck disable=SC2154
    if ! sudo "${sudoFlags[@]}" ufw allow "$ruleValue"; then
        msgbox "Failed to allow ${description}."
        return 1
    fi

    msgbox "${description} allowed."
    return 0
}

delete_rule() {
    local ruleValue="$1"
    local description="$2"

    # `sudoFlags` is initialized by the main launcher before sourcing this file.
    # shellcheck disable=SC2154
    if ! sudo "${sudoFlags[@]}" ufw delete allow "$ruleValue"; then
        msgbox "Failed to remove ${description}."
        return 1
    fi

    msgbox "${description} removed."
    return 0
}

allow_ssh_port() {
    local sshPort=""

    ensure_ufw || return 1
    ensure_firewall_backend_ready || return 1
    sshPort="$(resolve_ssh_port)" || {
        msgbox "Firewall change cancelled because the SSH port could not be confirmed."
        return 1
    }

    warn_sudo_authentication
    allow_rule "${sshPort}/tcp" "SSH port ${sshPort}/tcp"
}

enable_firewall() {
    ensure_ufw || return 1
    ensure_firewall_backend_ready || return 1
    allow_ssh_port || return 1

    # `sudoFlags` is initialized by the main launcher before sourcing this file.
    # shellcheck disable=SC2154
    if ! sudo "${sudoFlags[@]}" ufw --force enable; then
        msgbox "Failed to enable ufw."
        return 1
    fi

    # `sudoFlags` is initialized by the main launcher before sourcing this file.
    # shellcheck disable=SC2154
    if sudo "${sudoFlags[@]}" systemctl enable --now ufw; then
        msgbox "Firewall enabled."
        return 0
    fi

    msgbox "ufw rules were enabled, but the ufw service failed to enable at boot."
    return 1
}

disable_firewall() {
    ensure_ufw || return 1
    ensure_firewall_backend_ready || return 1

    # `sudoFlags` is initialized by the main launcher before sourcing this file.
    # shellcheck disable=SC2154
    if ! sudo "${sudoFlags[@]}" ufw --force disable; then
        msgbox "Failed to disable ufw."
        return 1
    fi

    # `sudoFlags` is initialized by the main launcher before sourcing this file.
    # shellcheck disable=SC2154
    if sudo "${sudoFlags[@]}" systemctl disable --now ufw; then
        msgbox "Firewall disabled."
        return 0
    fi

    msgbox "ufw rules were disabled, but the ufw service state could not be updated."
    return 1
}

open_custom_port() {
    local portNumber=""
    local protocolChoice=""

    ensure_ufw || return 1
    ensure_firewall_backend_ready || return 1
    portNumber="$(inputbox "Enter the port number to open.")" || return 1
    if ! valid_port "$portNumber"; then
        msgbox "Enter a valid port number from 1 to 65535."
        return 1
    fi

    protocolChoice="$(menulist "TCP" "UDP" "Both")" || return 1
    case "$protocolChoice" in
        "TCP")
            allow_rule "${portNumber}/tcp" "Port ${portNumber}/tcp"
            ;;
        "UDP")
            allow_rule "${portNumber}/udp" "Port ${portNumber}/udp"
            ;;
        "Both")
            allow_rule "${portNumber}/tcp" "Port ${portNumber}/tcp" || return 1
            allow_rule "${portNumber}/udp" "Port ${portNumber}/udp"
            ;;
    esac
}

list_simple_allow_rules() {
    local sshPort=""
    local statusLine=""
    declare -A tcpPorts=()
    declare -A udpPorts=()

    sshPort="$(resolve_ssh_port 2> /dev/null || true)"

    while IFS= read -r statusLine; do
        if [[ "$statusLine" =~ ^[[:space:]]*([0-9]+)/(tcp|udp)[[:space:]]+ALLOW[[:space:]]+Anywhere[[:space:]]*$ ]]; then
            if [[ "${BASH_REMATCH[1]}" == "$sshPort" ]]; then
                continue
            fi

            case "${BASH_REMATCH[2]}" in
                tcp)
                    tcpPorts["${BASH_REMATCH[1]}"]=1
                    ;;
                udp)
                    udpPorts["${BASH_REMATCH[1]}"]=1
                    ;;
            esac
        fi
    done < <(ufw_status_output)

    for portNumber in "${!tcpPorts[@]}"; do
        if [[ -n "${udpPorts[$portNumber]:-}" ]]; then
            printf '%s/both\n' "$portNumber"
            unset 'udpPorts[$portNumber]'
        else
            printf '%s/tcp\n' "$portNumber"
        fi
    done

    for portNumber in "${!udpPorts[@]}"; do
        printf '%s/udp\n' "$portNumber"
    done | sort -n -t / -k 1,1
}

close_port() {
    local removableRules=()
    local selectedRule=""
    local ruleValue=""

    ensure_ufw || return 1
    ensure_firewall_backend_ready || return 1
    while IFS= read -r ruleValue; do
        [[ -n "$ruleValue" ]] && removableRules+=("$ruleValue")
    done < <(list_simple_allow_rules | sort -n -t / -k 1,1)

    if [[ "${#removableRules[@]}" -eq 0 ]]; then
        msgbox "No simple removable port rules were found."
        return 0
    fi

    selectedRule="$(menulist "${removableRules[@]}")" || return 0
    case "$selectedRule" in
        */tcp|*/udp)
            delete_rule "$selectedRule" "Port ${selectedRule}"
            ;;
        */both)
            ruleValue="${selectedRule%/both}"
            delete_rule "${ruleValue}/tcp" "Port ${ruleValue}/tcp" || return 1
            delete_rule "${ruleValue}/udp" "Port ${ruleValue}/udp"
            ;;
    esac
}

view_firewall_status() {
    local tempFile=""
    local statusText=""

    ensure_ufw || return 1
    ensure_firewall_backend_ready || return 1
    tempFile="$(mktemp)"
    # `sudoFlags` is initialized by the main launcher before sourcing this file.
    # shellcheck disable=SC2154
    statusText="$(sudo "${sudoFlags[@]}" ufw status verbose 2>&1)"
    printf '%s\n' "$statusText" > "$tempFile"
    textbox "$tempFile"
    rm -f "$tempFile"
}

while true; do
    if firewall_enabled; then
        firewallToggleLabel="Disable firewall"
    else
        firewallToggleLabel="Enable firewall"
    fi

    firewallChoice="$(menulist \
        "$firewallToggleLabel" \
        "Allow SSH" \
        "Open port" \
        "Close port" \
        "View status" \
        "Back")" || break

    case "$firewallChoice" in
        "Enable firewall")
            enable_firewall
            ;;
        "Disable firewall")
            disable_firewall
            ;;
        "Allow SSH")
            allow_ssh_port
            ;;
        "Open port")
            open_custom_port
            ;;
        "Close port")
            close_port
            ;;
        "View status")
            view_firewall_status
            ;;
        "Back")
            break
            ;;
    esac
done