#!/usr/bin/bash
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.  Please see LICENSE at the top level of
# the source code distribution for details.
#
# @package nftables-apply
# @author <felix.bouynot@setenforce.one>
# @link https://codeberg.org/fbouynot/nftables-apply/
# @copyright <felix.bouynot@setenforce.one>
#
# Free adaptation for nftables of iptables-apply (https://github.com/wertarbyte/iptables/blob/master/iptables-apply)
#

# -e: When a command fails, bash exits instead of continuing with the rest of the script
# -u: This will make the script fail, when accessing an unset variable
# -o pipefail: This will ensure that a pipeline command is treated as failed, even if one command in the pipeline fails
set -euo pipefail

# Replace the Internal Field Separator ' \n\t' by '\n\t' so you can loop through names with spaces
IFS=$'\n\t'

# Enable debug mode by running your script as TRACE=1 ./script.sh instead of ./script.sh
if [[ "${TRACE-0}" == "1" ]]
then
    set -o xtrace
fi

# Define constants
PROGNAME="${0##*/}"
VERSION='2.1.0'
RED="$(tput setaf 1)"
NC="$(tput sgr0)" # No Color

DEFAULT_TIMEOUT=15
DEFAULT_DESTINATION_FILE='/etc/nftables/main.nft'
DEFAULT_SOURCE_FILE='/etc/nftables/candidate.nft'

readonly PROGNAME VERSION RED NC DEFAULT_TIMEOUT DEFAULT_DESTINATION_FILE DEFAULT_SOURCE_FILE

cleanup() {
    if [[ -f "${BACKUP_FILE}" ]]
    then
        printf '%sE:%s backup file preserved at %s\n' "${RED}" "${NC}" "${BACKUP_FILE}" >&2
    else
        rm -rf "${BACKUP_DIR}"
    fi
}

help() {
    cat << EOF
Usage: ${PROGNAME} [-Vh] [ { -s | --source-file } <source-file> ] [ { -d | --destination-file } <destination-file> ] [ { -t | --timeout } <timeout> ]
-h    --help                                                     Print this message.
-V    --version                                                  Print the version.
-s    --source-file        STRING                                The source file for candidate config.           (default: ${DEFAULT_SOURCE_FILE})
-d    --destination-file   STRING                                The destination file where to write the config. (default: ${DEFAULT_DESTINATION_FILE})
-t    --timeout            INT                                   The time to wait before rolling back.           (default: ${DEFAULT_TIMEOUT})
EOF

    exit "${1:-0}"
}

version() {
    cat << EOF
${PROGNAME} version ${VERSION} under GPLv3 licence.
EOF

    exit 0
}

# Deal with arguments
while [[ $# -gt 0 ]]
do
    key="${1}"

    case $key in
        -h|--help)
            help
            ;;
        -s|--source-file)
            if [[ $# -lt 2 ]]; then
                printf '%sE:%s %s requires a value\n' "${RED}" "${NC}" "${1}" >&2
                help 2
            fi
            source_file="${2}"
            shift # consume -s
            ;;
        -d|--destination-file)
            if [[ $# -lt 2 ]]; then
                printf '%sE:%s %s requires a value\n' "${RED}" "${NC}" "${1}" >&2
                help 2
            fi
            destination_file="${2}"
            shift # consume -d
            ;;
        -t|--timeout)
            if [[ $# -lt 2 ]]; then
                printf '%sE:%s %s requires a value\n' "${RED}" "${NC}" "${1}" >&2
                help 2
            fi
            rollback_timeout="${2}"
            shift # consume -t
            ;;
        -V|--version)
            version
            ;;
        *)
            printf '%sE:%s unknown option: %s\n' "${RED}" "${NC}" "${1}" >&2
            help 2
            ;;
    esac
    shift # consume $1
done

# Set defaults if no options specified
source_file="${source_file:-$DEFAULT_SOURCE_FILE}"
destination_file="${destination_file:-$DEFAULT_DESTINATION_FILE}"
rollback_timeout="${rollback_timeout:-$DEFAULT_TIMEOUT}"

# =~ matches a regexp: ^[1-9][0-9]*$ accepts only strictly positive integers
if [[ ! "${rollback_timeout}" =~ ^[1-9][0-9]*$ ]]
then
    printf '%sE:%s timeout must be a positive integer\n' "${RED}" "${NC}" >&2
    exit 9
fi

readonly source_file destination_file rollback_timeout

# Check root permissions
check_root() {
    # Check the command is run as root
    if [ "${EUID}" -ne 0 ]
    then
        printf '%sE:%s please run as root\n' "${RED}" "${NC}" >&2
        exit 3
    fi
}

stop_fail2ban() {
    if systemctl is-active fail2ban > /dev/null 2>&1
    then
        systemctl stop fail2ban 2>/dev/null
    fi
}

start_fail2ban() {
    if systemctl is-enabled fail2ban > /dev/null 2>&1
    then
        systemctl start fail2ban 2>/dev/null
    fi
}

restore() {
    nft flush ruleset
    nft -f "${BACKUP_FILE}"
    rm -f "${BACKUP_FILE}"
    start_fail2ban
}

save() {
    cp "${source_file}" "${destination_file}"
    printf '\nConfiguration changed\n'
}

# Main function
main() {
    # Check the command is run as root
    check_root

    # Check if we can read the destination file
    if [[ ! -r "${destination_file}" ]]
    then
        printf '%sE:%s cannot read %s\n' "${RED}" "${NC}" "${destination_file}" >&2
        exit 4
    fi

    # Check if we can read the source file
    if [[ ! -r "${source_file}" ]]
    then
        printf '%sE:%s cannot read %s\n' "${RED}" "${NC}" "${source_file}" >&2
        exit 5
    fi

    # Backup current ruleset
    BACKUP_DIR="$(mktemp -d /tmp/nftables-apply.XXXXXXXXXX)"
    BACKUP_FILE="${BACKUP_DIR}/nftables.conf.bak"
    trap cleanup EXIT
    nft list ruleset > "${BACKUP_FILE}"

    # Dry run new ruleset, exit if failures
    nft --check -f "${source_file}" || { printf '%sE:%s Invalid rules, exiting\n' "${RED}" "${NC}" >&2; exit 6; }

    # Check the candidate configuration starts by flushing ruleset
    if [[ $(head -n 1 "${source_file}") != "flush ruleset" ]]
    then
        sed -i '1s/^/flush ruleset\n/' "${source_file}"
    fi

    stop_fail2ban

    # Apply new ruleset, rollback if timeout
    timeout "${rollback_timeout}s" nft -f "${source_file}" || { printf '%sE:%s timeout while applying new configuration, rolling back to the previous ruleset\n' "${RED}" "${NC}" >&2; restore; exit 7; }

    # Ask the user if they can open a new connection
    # If they can't, rollback
    # If they can, save
    printf 'Can you establish NEW connections to the machine? (y/N) '
    read -r -n1 -t "${rollback_timeout}" answer 2>&1 || :
    if [[ "${answer}" == "y" ]]
    then
        save
    else
        printf '\n%sE:%s rolling back to the previous ruleset\n' "${RED}" "${NC}" >&2
        restore
        exit 8
    fi
    rm -f "${BACKUP_FILE}"
    start_fail2ban

    exit 0
}

main "$@"
