#!/bin/bash
#
# legacy_iptables_rules.sh
#
# Copyright (C) 2023, Branko Majic <branko@majic.rs>
#
# 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.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

# Treat unset variables as errors.
set -u

PROGRAM="legacy_iptables_rules.sh"

function usage() {
    cat <<EOF
$PROGRAM, helper tool for removing legacy iptables rules

Usage:
  $PROGRAM [OPTIONS] remove
EOF
}

function short_help() {
    cat <<EOF
$(usage)

For more details see $PROGRAM -h.
EOF
}

function long_help() {
    cat <<EOF
$(usage)

$PROGRAM is a helper tool that can be used to remove the legacy
iptables rules.

The tool works by resetting the default policies on all the relevant
chains and tables, flushing the rules, as well as unloading the
related kernel modules.

Tool implements multiple commands, as documented below.

  remove

    Removes the legacy iptables rules.

$PROGRAM accepts the following options:

    -q
        Quiet mode.
    -d
        Enable debug mode.
    -v
        Show script version and licensing information.
    -h
        Show full help.

Please report bugs and send feature requests to <branko@majic.rs>.
EOF
}

function version() {
    cat <<EOF
$PROGRAM

+-----------------------------------------------------------------------+
| Copyright (C) 2023, Branko Majic <branko@majic.rs>                    |
|                                                                       |
| 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.                                   |
|                                                                       |
| This program is distributed in the hope that it will be useful,       |
| but WITHOUT ANY WARRANTY; without even the implied warranty of        |
| MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         |
| GNU General Public License for more details.                          |
|                                                                       |
| You should have received a copy of the GNU General Public License     |
| along with this program.  If not, see <http://www.gnu.org/licenses/>. |
+-----------------------------------------------------------------------+

EOF
}

# Set-up colours for message printing if we're not piping and terminal is
# capable of outputting the colors.
_COLOR_TERMINAL=$(tput colors 2>&1)
if [[ -t 1 ]] && (( _COLOR_TERMINAL > 0 )); then
    _TEXT_BOLD=$(tput bold)
    _TEXT_WHITE=$(tput setaf 7)
    _TEXT_BLUE=$(tput setaf 6)
    _TEXT_GREEN=$(tput setaf 2)
    _TEXT_YELLOW=$(tput setaf 3)
    _TEXT_RED=$(tput setaf 1)
    _TEXT_RESET=$(tput sgr0)
else
    _TEXT_BOLD=""
    _TEXT_WHITE=""
    _TEXT_BLUE=""
    _TEXT_GREEN=""
    _TEXT_YELLOW=""
    _TEXT_RED=""
    _TEXT_RESET=""
fi

# Set-up functions for printing coloured messages.
function debug() {
    if [[ $DEBUG != 0 ]]; then
        echo "${_TEXT_BOLD}${_TEXT_BLUE}[DEBUG]${_TEXT_RESET}" "$@"
    fi
}

function info() {
    if [[ $QUIET == 0 ]]; then
        echo "${_TEXT_BOLD}${_TEXT_WHITE}[INFO] ${_TEXT_RESET}" "$@"
    fi
}

function success() {
    if [[ $QUIET == 0 ]]; then
        echo "${_TEXT_BOLD}${_TEXT_GREEN}[OK]   ${_TEXT_RESET}" "$@"
    fi
}

function warning() {
    echo "${_TEXT_BOLD}${_TEXT_YELLOW}[WARN] ${_TEXT_RESET}" "$@" >&2
}

function error() {
    echo "${_TEXT_BOLD}${_TEXT_RED}[ERROR]${_TEXT_RESET}" "$@" >&2
}

#
# Removes legacy iptables (both IPv4 and IPv6) rules.
#
# Returns:
#
#   0 in case of success, 1 otherwise.
#
function remove_legacy_iptables() {
    local family table tables chain iptables_save iptables_legacy module
    declare -a removed_families=()

    family=(ip ip6)
    tables=(filter nat mangle raw security)

    for family in "${family[@]}"; do

        iptables_save="${family}tables-save"
        iptables_legacy="${family}tables-legacy"
        iptables_legacy_save="${family}tables-legacy-save"

        if "$iptables_save" 2>&1 | grep -q "Warning: ${family}tables-legacy tables present"; then
            info "Legacy iptables detected for IP family: $family"

            for table in "${tables[@]}"; do

                info "Processing table: $family/$table"

                # Iterate over all chains with default policy.
                for chain in $("$iptables_legacy_save" -t "$table" | grep -E '^:.*(ACCEPT|DROP)' | sed -e 's/^://;s/ .*//'); do

                    info "Setting default policy for chain: $family/$table/$chain"
                    if ! "$iptables_legacy" -t "$table" -P "$chain" ACCEPT; then
                        error "Failed to set policy for chain: $chain ($table)"
                        return 1
                    fi

                done

                info "Flushing table: $family/$table"
                if ! "$iptables_legacy" -t "$table" -F; then
                    error "Failed to flush the table: $table"
                    return 1
                fi

                info "Removing user chains: $family/$table"
                if ! "$iptables_legacy" -t "$table" -X; then
                    error "Failed to remove user chains: $family/$table"
                    return 1
                fi

            done

            # Modules are named after a table.
            for table in "${tables[@]}"; do
                module="${family}table_${table}"

                info "Unloading module: $module"
                if ! rmmod "$module"; then
                    error "Failed to unload module: $module"
                    return 1
                fi
            done

            removed_families+=("$family")

        fi

    done

    if (( ${#removed_families[@]} > 0 )); then
        success "Removed legacy iptables for families: ${removed_families[*]}"
    else
        success "No legacy iptables are present on the system."
    fi

    return 0
}

# Define error codes.
SUCCESS=0
ERROR_ARGUMENTS=1
ERROR_GENERAL=2

# Disable debug and quiet modes by default.
DEBUG=0
QUIET=0

# If no arguments were given, just show usage help.
if [[ -z ${1-} ]]; then
    short_help
    exit "$ERROR_ARGUMENTS"
fi

# Parse the arguments
while getopts "qdvh" opt; do
    case "$opt" in
	q) QUIET=1;;
	d) DEBUG=1;;
        v) version
           exit "$SUCCESS";;
        h) long_help
           exit "$SUCCESS";;
        *) short_help
           exit "$ERROR_ARGUMENTS";;
    esac
done
i=$OPTIND
shift $(( i-1 ))

# Quiet and debug are mutually exclusive.
if [[ $QUIET != 0 && $DEBUG != 0 ]]; then
    error "Quiet and debug options are mutually exclusive."
    exit "$ERROR_ARGUMENTS"
fi

COMMAND="${1-}"

if [[ $COMMAND == remove ]]; then

    if ! remove_legacy_iptables; then
        error "Failed to remove legacy iptables."
        exit "$ERROR_GENERAL"
    fi

else

    error "Unsupported command: $COMMAND"
    exit "$ERROR_ARGUMENTS"

fi