#!/bin/bash

# Author: Jose Mendez
# Blog site: sandbox99.cc
# This script sets default policies to DROP and adds rules
# to allow necessary traffic (SSH, HTTP, HTTPS, localhost).
# It also adds LOG rules for monitoring traffic.
# It checks if policies/rules are already configured before applying.

# Require root privileges to run the script
if [[ $EUID -ne 0 ]]; then
   echo "This script must be run as root."
   exit 1
fi

echo "Starting iptables hardening setup with logging..."

# --- Configure Default Policies ---

# Desired policies
IPV4_INPUT_POLICY="DROP"
IPV4_FORWARD_POLICY="DROP"
IPV4_OUTPUT_POLICY="ACCEPT"

IPV6_INPUT_POLICY="DROP"
IPV6_FORWARD_POLICY="DROP"
IPV6_OUTPUT_POLICY="ACCEPT"

# Check and set IPv4 policies
echo "Checking IPv4 default policies..."
# Use iptables-save and parse for robustness
current_ipv4_policies=$(sudo iptables-save 2>/dev/null | grep ":INPUT\|:FORWARD\|:OUTPUT")
current_input_policy=$(echo "$current_ipv4_policies" | grep ":INPUT" | awk '{print $NF}')
current_forward_policy=$(echo "$current_ipv4_policies" | grep ":FORWARD" | awk '{print $NF}')
current_output_policy=$(echo "$current_ipv4_policies" | grep ":OUTPUT" | awk '{print $NF}')


if [ "$current_input_policy" != "$IPV4_INPUT_POLICY" ]; then
    echo "Setting IPv4 INPUT policy to $IPV4_INPUT_POLICY..."
    iptables -P INPUT "$IPV4_INPUT_POLICY"
else
    echo "IPv4 INPUT policy is already set to $IPV4_INPUT_POLICY."
fi

if [ "$current_forward_policy" != "$IPV4_FORWARD_POLICY" ]; then
    echo "Setting IPv4 FORWARD policy to $IPV4_FORWARD_POLICY..."
    iptables -P FORWARD "$IPV4_FORWARD_POLICY"
else
    echo "IPv4 FORWARD policy is already set to $IPV4_FORWARD_POLICY."
fi

if [ "$current_output_policy" != "$IPV4_OUTPUT_POLICY" ]; then
    echo "Setting IPv4 OUTPUT policy to $IPV4_OUTPUT_POLICY..."
    iptables -P OUTPUT "$IPV4_OUTPUT_POLICY"
else
    echo "IPv4 OUTPUT policy is already set to $IPV4_OUTPUT_POLICY."
fi

# Check and set IPv6 policies
echo "Checking IPv6 default policies..."
if command -v ip6tables &> /dev/null; then
    current_ipv6_policies=$(sudo ip6tables-save 2>/dev/null | grep ":INPUT\|:FORWARD\|:OUTPUT")
    current_input_policy_v6=$(echo "$current_ipv6_policies" | grep ":INPUT" | awk '{print $NF}')
    current_forward_policy_v6=$(echo "$current_ipv6_policies" | grep ":FORWARD" | awk '{print $NF}')
    current_output_policy_v6=$(echo "$current_ipv6_policies" | grep ":OUTPUT" | awk '{print $NF}')

    if [ "$current_input_policy_v6" != "$IPV6_INPUT_POLICY" ]; then
        echo "Setting IPv6 INPUT policy to $IPV6_INPUT_POLICY..."
        ip6tables -P INPUT "$IPV6_INPUT_POLICY"
    else
        echo "IPv6 INPUT policy is already set to $IPV6_INPUT_POLICY."
    fi

    if [ "$current_forward_policy_v6" != "$IPV6_FORWARD_POLICY" ]; then
        echo "Setting IPv6 FORWARD policy to $IPV6_FORWARD_POLICY..."
        ip6tables -P FORWARD "$IPV6_FORWARD_POLICY"
    else
        echo "IPv6 FORWARD policy is already set to $IPV6_FORWARD_POLICY."
    fi

    if [ "$current_output_policy_v6" != "$IPV6_OUTPUT_POLICY" ]; then
        echo "Setting IPv6 OUTPUT policy to $IPV6_OUTPUT_POLICY..."
        ip6tables -P OUTPUT "$IPV6_OUTPUT_POLICY"
    else
        echo "IPv6 OUTPUT policy is already set to $IPV6_OUTPUT_POLICY."
    fi
else
    echo "ip6tables command not found or saved rules not found. Skipping IPv6 policy configuration."
fi


# --- Add Essential Rules (Check if they exist first) ---

# Rule: Allow established and related incoming connections (IPv4)
echo "Checking for IPv4 established/related rule..."
if ! iptables -C INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT 2>/dev/null; then
    echo "Adding IPv4 established/related rule..."
    iptables -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
else
    echo "IPv4 established/related rule already exists."
fi

# Rule: Allow established and related incoming connections (IPv6)
echo "Checking for IPv6 established/related rule..."
if command -v ip6tables &> /dev/null; then
    if ! ip6tables -C INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT 2>/dev/null; then
        echo "Adding IPv6 established/related rule..."
        ip6tables -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    else
        echo "IPv6 established/related rule already exists."
    fi
else
    echo "ip6tables command not found. Skipping IPv6 established/related rule."
fi

# Rule: Allow SSH incoming connections (Default port 22) (IPv4)
# Modify the port if your SSH server uses a different one
SSH_PORT="22"
echo "Checking for IPv4 SSH rule on port $SSH_PORT..."
if ! iptables -C INPUT -p tcp --dport "$SSH_PORT" -j ACCEPT 2>/dev/null; then
    echo "Adding IPv4 SSH rule on port $SSH_PORT..."
    iptables -A INPUT -p tcp --dport "$SSH_PORT" -j ACCEPT
else
    echo "IPv4 SSH rule on port $SSH_PORT already exists."
fi

# Rule: Allow SSH incoming connections (Default port 22) (IPv6)
echo "Checking for IPv6 SSH rule on port $SSH_PORT..."
if command -v ip6tables &> /dev/null; then
    if ! ip6tables -C INPUT -p tcp --dport "$SSH_PORT" -j ACCEPT 2>/dev/null; then
        echo "Adding IPv6 SSH rule on port $SSH_PORT..."
        ip6tables -A INPUT -p tcp --dport "$SSH_PORT" -j ACCEPT
    else
        echo "IPv6 SSH rule on port $SSH_PORT already exists."
    fi
else
    echo "ip6tables command not found. Skipping IPv6 SSH rule."
fi

# Rule: Allow HTTP incoming connections (Port 80) (IPv4)
HTTP_PORT="80"
echo "Checking for IPv4 HTTP rule on port $HTTP_PORT..."
if ! iptables -C INPUT -p tcp --dport "$HTTP_PORT" -j ACCEPT 2>/dev/null; then
    echo "Adding IPv4 HTTP rule on port $HTTP_PORT..."
    iptables -A INPUT -p tcp --dport "$HTTP_PORT" -j ACCEPT
else
    echo "IPv4 HTTP rule on port $HTTP_PORT already exists."
fi

# Rule: Allow HTTP incoming connections (Port 80) (IPv6)
echo "Checking for IPv6 HTTP rule on port $HTTP_PORT..."
if command -v ip6tables &> /dev/null; then
    if ! ip6tables -C INPUT -p tcp --dport "$HTTP_PORT" -j ACCEPT 2>/dev/null; then
        echo "Adding IPv6 HTTP rule on port "$HTTP_PORT"..."
        ip6tables -A INPUT -p tcp --dport "$HTTP_PORT" -j ACCEPT
    else
        echo "IPv6 HTTP rule on port "$HTTP_PORT" already exists."
    fi
else
    echo "ip6tables command not found. Skipping IPv6 HTTP rule."
fi

# Rule: Allow HTTPS incoming connections (Port 443) (IPv4)
HTTPS_PORT="443"
echo "Checking for IPv4 HTTPS rule on port $HTTPS_PORT..."
if ! iptables -C INPUT -p tcp --dport "$HTTPS_PORT" -j ACCEPT 2>/dev/null; then
    echo "Adding IPv4 HTTPS rule on port "$HTTPS_PORT"..."
    iptables -A INPUT -p tcp --dport "$HTTPS_PORT" -j ACCEPT
else
    echo "IPv4 HTTPS rule on port "$HTTPS_PORT" already exists."
fi

# Rule: Allow HTTPS incoming connections (Port 443) (IPv6)
echo "Checking for IPv6 HTTPS rule on port $HTTPS_PORT..."
if command -v ip6tables &> /dev/null; then
    if ! ip6tables -C INPUT -p tcp --dport "$HTTPS_PORT" -j ACCEPT 2>/dev/null; then
        echo "Adding IPv6 HTTPS rule on port "$HTTPS_PORT"..."
        ip6tables -A INPUT -p tcp --dport "$HTTPS_PORT" -j ACCEPT
    else
        echo "IPv6 HTTPS rule on port "$HTTPS_PORT" already exists."
    fi
else
    echo "ip6tables command not found. Skipping IPv6 HTTPS rule."
fi


# Rule: Allow localhost traffic (IPv4)
echo "Checking for IPv4 localhost rules..."
if ! iptables -C INPUT -i lo -j ACCEPT 2>/dev/null; then
    echo "Adding IPv4 localhost INPUT rule..."
    iptables -A INPUT -i lo -j ACCEPT
else
    echo "IPv4 localhost INPUT rule already exists."
fi
if ! iptables -C OUTPUT -o lo -j ACCEPT 2>/dev/null; then
     echo "Adding IPv4 localhost OUTPUT rule..."
     iptables -A OUTPUT -o lo -j ACCEPT
else
     echo "IPv4 localhost OUTPUT rule already exists."
fi


# Rule: Allow localhost traffic (IPv6)
echo "Checking for IPv6 localhost rules..."
if command -v ip6tables &> /dev/null; then
    if ! ip6tables -C INPUT -i lo -j ACCEPT 2>/dev/null; then
        echo "Adding IPv6 localhost INPUT rule..."
        ip6tables -A INPUT -i lo -j ACCEPT
    else
        echo "IPv6 localhost INPUT rule already exists."
    fi
    if ! ip6tables -C OUTPUT -o lo -j ACCEPT 2>/dev/null; then
         echo "Adding IPv6 localhost OUTPUT rule..."
         ip6tables -A OUTPUT -o lo -j ACCEPT
    else
         echo "IPv6 localhost OUTPUT rule already exists."
    fi
else
    echo "ip6tables command not found. Skipping IPv6 localhost rules."
fi


# --- Add LOG Rules (Before default DROP) ---
# These rules log packets that would otherwise be dropped by the default policy
# Use a rate limit to prevent log flooding.

LOG_LIMIT="5/min" # Limit to 5 log messages per minute per rule

# Log IPv4 packets dropped by default INPUT policy
LOG_PREFIX_INPUT_DROP="iptables-input-dropped: "
echo "Checking for IPv4 INPUT drop logging rule..."
if ! iptables -C INPUT -m limit --limit "$LOG_LIMIT" -j LOG --log-prefix "$LOG_PREFIX_INPUT_DROP" 2>/dev/null; then
    echo "Adding IPv4 INPUT drop logging rule..."
    # Place this rule near the end of the chain, before the default policy
    # Check if there's a jump to a different chain near the end and place before it if needed.
    # For simplicity, we'll append it. If you have complex jumps, verify placement manually.
    iptables -A INPUT -m limit --limit "$LOG_LIMIT" -j LOG --log-prefix "$LOG_PREFIX_INPUT_DROP"
else
    echo "IPv4 INPUT drop logging rule already exists."
fi

# Log IPv6 packets dropped by default INPUT policy
LOG_PREFIX_INPUT_DROP_V6="ip6tables-input-dropped: "
echo "Checking for IPv6 INPUT drop logging rule..."
if command -v ip6tables &> /dev/null; then
    if ! ip6tables -C INPUT -m limit --limit "$LOG_LIMIT" -j LOG --log-prefix "$LOG_PREFIX_INPUT_DROP_V6" 2>/dev/null; then
        echo "Adding IPv6 INPUT drop logging rule..."
        ip6tables -A INPUT -m limit --limit "$LOG_LIMIT" -j LOG --log-prefix "$LOG_PREFIX_INPUT_DROP_V6"
    else
        echo "IPv6 INPUT drop logging rule already exists."
    fi
else
    echo "ip6tables command not found. Skipping IPv6 INPUT drop logging rule."
fi


# Log IPv4 packets dropped by default FORWARD policy
LOG_PREFIX_FORWARD_DROP="iptables-forward-dropped: "
echo "Checking for IPv4 FORWARD drop logging rule..."
if ! iptables -C FORWARD -m limit --limit "$LOG_LIMIT" -j LOG --log-prefix "$LOG_PREFIX_FORWARD_DROP" 2>/dev/null; then
    echo "Adding IPv4 FORWARD drop logging rule..."
    # Place this rule near the end of the chain, before the default policy
    iptables -A FORWARD -m limit --limit "$LOG_LIMIT" -j LOG --log-prefix "$LOG_PREFIX_FORWARD_DROP"
else
    echo "IPv4 FORWARD drop logging rule already exists."
fi

# Log IPv6 packets dropped by default FORWARD policy
LOG_PREFIX_FORWARD_DROP_V6="ip6tables-forward-dropped: "
echo "Checking for IPv6 FORWARD drop logging rule..."
if command -v ip6tables &> /dev/null; then
    if ! ip6tables -C FORWARD -m limit --limit "$LOG_LIMIT" -j LOG --log-prefix "$LOG_PREFIX_FORWARD_DROP_V6" 2>/dev/null; then
        echo "Adding IPv6 FORWARD drop logging rule..."
        ip6tables -A FORWARD -m limit --limit "$LOG_LIMIT" -j LOG --log-prefix "$LOG_PREFIX_FORWARD_DROP_V6"
    else
        echo "IPv6 FORWARD drop logging rule already exists."
    fi
else
    echo "ip6tables command not found. Skipping IPv6 FORWARD drop logging rule."
fi

# Log IPv4 packets hitting the DOCKER-USER chain (Optional, for debugging Docker traffic)
# Place this rule strategically in DOCKER-USER depending on what you want to log.
# Appending it logs packets that weren't handled by specific rules before it.
LOG_PREFIX_DOCKER_USER="iptables-docker-user: "
echo "Checking for IPv4 DOCKER-USER logging rule..."
if ! iptables -C DOCKER-USER -m limit --limit "$LOG_LIMIT" -j LOG --log-prefix "$LOG_PREFIX_DOCKER_USER" 2>/dev/null; then
    echo "Adding IPv4 DOCKER-USER logging rule..."
     # It's often useful to log traffic *entering* the DOCKER-USER chain, which would
     # require inserting the rule at the beginning (-I DOCKER-USER 1).
     # For simplicity here, we append it. Adjust manually if needed.
    iptables -A DOCKER-USER -m limit --limit "$LOG_LIMIT" -j LOG --log-prefix "$LOG_PREFIX_DOCKER_USER"
else
    echo "IPv4 DOCKER-USER logging rule already exists."
fi

# Log IPv6 packets hitting the DOCKER-USER chain (Optional)
LOG_PREFIX_DOCKER_USER_V6="ip6tables-docker-user: "
echo "Checking for IPv6 DOCKER-USER logging rule..."
if command -v ip6tables &> /dev/null; then
    if ! ip6tables -C DOCKER-USER -m limit --limit "$LOG_LIMIT" -j LOG --log-prefix "$LOG_PREFIX_DOCKER_USER_V6" 2>/dev/null; then
        echo "Adding IPv6 DOCKER-USER logging rule..."
        ip6tables -A DOCKER-USER -m limit --limit "$LOG_LIMIT" -j LOG --log-prefix "$LOG_PREFIX_DOCKER_USER_V6"
    else
        echo "IPv6 DOCKER-USER logging rule already exists."
    fi
else
    echo "ip6tables command not found. Skipping IPv6 DOCKER-USER logging rule."
fi


# --- Recommended Rule for Docker Integration (DOCKER-USER Chain) ---
# This rule sends traffic from the DOCKER-USER chain to the RETURN target.
# Place your custom DOCKER-USER filtering rules *before* this RETURN rule.

echo "Checking for basic IPv4 DOCKER-USER chain rule (RETURN)..."
if ! iptables -C DOCKER-USER -j RETURN 2>/dev/null; then
    echo "Adding basic IPv4 DOCKER-USER chain rule (RETURN)..."
    iptables -A DOCKER-USER -j RETURN
else
    echo "Basic IPv4 DOCKER-USER chain rule (RETURN) already exists."
fi

echo "Checking for basic IPv6 DOCKER-USER chain rule (RETURN)..."
if command -v ip6tables &> /dev/null; then
    if ! ip6tables -C DOCKER-USER -j RETURN 2>/dev/null; then
        echo "Adding basic IPv6 DOCKER-USER chain rule (RETURN)..."
        ip6tables -A DOCKER-USER -j RETURN
    else
        echo "Basic IPv6 DOCKER-USER chain rule (RETURN) already exists."
    fi
else
    echo "ip6tables command not found. Skipping basic IPv6 DOCKER-USER rule."
fi


# --- Save Rules ---
echo "Saving iptables rules..."
if command -v netfilter-persistent &> /dev/null; then
    netfilter-persistent save
    echo "iptables rules saved using netfilter-persistent."
else
    echo "netfilter-persistent not found. Please install it ('sudo apt install iptables-persistent') to make rules persistent across reboots."
    echo "You can manually save rules using: sudo iptables-save > /etc/iptables/rules.v4"
    if command -v ip6tables &> /dev/null; then
        echo "And IPv6 rules using: sudo ip6tables-save > /etc/iptables/rules.v6"
    fi
fi

echo "iptables hardening setup with logging complete."
echo "Monitor logs using 'sudo journalctl -k -f | grep \"iptables-\"' or by checking /var/log/kern.log or /var/log/syslog."
echo "Remember to add any specific filtering rules for Docker containers in the DOCKER-USER chain before the RETURN rule."