#!/bin/bash
# Author:  ServerStack
#
# Notes: Standalone iptables firewall installer, extracted from
#        core/os/init_Ubuntu.sh / core/os/init_Debian.sh.
#        Only Debian / Ubuntu are supported (iptables-persistent).
#
# Usage:
#   bash install_iptables.sh                 # auto-detect SSH port
#   bash install_iptables.sh --ssh-port 2222 # use a custom SSH port
#   SSH_PORT=2222 bash install_iptables.sh   # alternative: env var
#
# It installs iptables-persistent non-interactively, writes a default
# ruleset (INPUT DROP; allow lo / RELATED,ESTABLISHED / tcp 22,80,443 / icmp),
# and mirrors an IPv6 variant. Idempotent: an existing valid ruleset is kept.

# --- minimal color helpers (self-contained, no lib/color.sh needed) ---
CSI=$'\033['
CEND="${CSI}0m"
CSUCCESS="${CSI}32m"
CFAILURE="${CSI}1;31m"
CWARNING="${CSI}1;33m"
CMSG="${CSI}1;36m"

# --- usage / help ---
usage() {
  cat << EOF
${CMSG}Standalone iptables firewall installer${CEND}
Only Debian / Ubuntu are supported (iptables-persistent).

${CMSG}Usage:${CEND}
  bash install_iptables.sh                 # auto-detect SSH port (from /etc/ssh/sshd_config, default 22)
  bash install_iptables.sh --ssh-port 2222 # use a custom SSH port (2222)
  SSH_PORT=2222 bash install_iptables.sh   # alternative: set SSH port via env var

${CMSG}Options:${CEND}
  --ssh-port <N>   Allow inbound TCP on SSH port N instead of the detected one (22).
  -h, --help       Show this help message and exit.

It installs iptables-persistent non-interactively, writes a default ruleset
(INPUT DROP; allow lo / RELATED,ESTABLISHED / tcp 22,80,443 / icmp), and mirrors
an IPv6 variant. Idempotent: an existing valid ruleset is kept.
EOF
}

# --- argument parsing ---
ssh_port=""
while [ $# -gt 0 ]; do
  case "$1" in
    --ssh-port)
      ssh_port="$2"
      shift 2
      ;;
    --ssh-port=*)
      ssh_port="${1#*=}"
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "${CWARNING}Unknown argument: $1${CEND}" >&2
      usage
      exit 1
      ;;
  esac
done

[ -z "${ssh_port}" ] && ssh_port="${SSH_PORT:-}"

# --- detect the current SSH port from sshd_config ---
if [ -z "${ssh_port}" ]; then
  if [ -e "/etc/ssh/sshd_config" ] && [ -n "$(grep '^Port ' /etc/ssh/sshd_config)" ]; then
    ssh_port="$(grep '^Port ' /etc/ssh/sshd_config | awk '{print $2}' | head -1)"
  else
    ssh_port=22
  fi
fi

echo "${CMSG}Installing iptables firewall tooling (SSH port: ${ssh_port})...${CEND}"

# --- package install (non-interactive via debconf) ---
export DEBIAN_FRONTEND=noninteractive
apt-get -y install debconf-utils
echo iptables-persistent iptables-persistent/autosave_v4 boolean true | debconf-set-selections
echo iptables-persistent iptables-persistent/autosave_v6 boolean true | debconf-set-selections
apt-get -y install iptables-persistent

# --- only write a fresh ruleset if the existing one is not valid ---
IPTABLES_STATUS=no
if [ -e "/etc/iptables/rules.v4" ] \
   && [ -n "$(grep '^:INPUT DROP' /etc/iptables/rules.v4)" ] \
   && [ -n "$(grep 'NEW -m tcp --dport 22 -j ACCEPT' /etc/iptables/rules.v4)" ] \
   && [ -n "$(grep 'NEW -m tcp --dport 80 -j ACCEPT' /etc/iptables/rules.v4)" ]; then
  IPTABLES_STATUS=yes
fi

if [ "${IPTABLES_STATUS}" == "no" ]; then
  cat > /etc/iptables/rules.v4 << EOF
# Firewall configuration written by system-config-securitylevel
# Manual customization of this file is not recommended.
*filter
:INPUT DROP [0:0]
:FORWARD ACCEPT [0:0]
:OUTPUT ACCEPT [0:0]
:syn-flood - [0:0]
-A INPUT -i lo -j ACCEPT
-A INPUT -m state --state RELATED,ESTABLISHED -j ACCEPT
-A INPUT -p tcp -m state --state NEW -m tcp --dport 22 -j ACCEPT
-A INPUT -p tcp -m state --state NEW -m tcp --dport 80 -j ACCEPT
-A INPUT -p tcp -m state --state NEW -m tcp --dport 443 -j ACCEPT
-A INPUT -p icmp -m icmp --icmp-type 8 -j ACCEPT
COMMIT
EOF
fi

# --- add a rule for a non-default SSH port ---
FW_PORT_FLAG=$(grep -ow "dport ${ssh_port}" /etc/iptables/rules.v4)
if [ -z "${FW_PORT_FLAG}" ] && [ "${ssh_port}" != "22" ]; then
  sed -i "s@dport 22 -j ACCEPT@&\n-A INPUT -p tcp -m state --state NEW -m tcp --dport ${ssh_port} -j ACCEPT@" /etc/iptables/rules.v4
fi

# --- apply IPv4, then mirror IPv6 ---
iptables-restore < /etc/iptables/rules.v4
/bin/cp /etc/iptables/rules.v{4,6}
sed -i 's@icmp@icmpv6@g' /etc/iptables/rules.v6
ip6tables-restore < /etc/iptables/rules.v6
ip6tables-save > /etc/iptables/rules.v6

if [ -e "/etc/iptables/rules.v4" ]; then
  echo; echo "${CSUCCESS}iptables firewall installed and applied successfully!${CEND}"
  echo "${CMSG}Current INPUT rules:${CEND}"
  iptables -L INPUT -n --line-numbers
else
  echo; echo "${CFAILURE}iptables install failed, Please try again!${CEND}"
  exit 1
fi