#!/bin/bash
set -euxo pipefail

export DEBIAN_FRONTEND=noninteractive

# --- Base updates and security ---
apt-get update
apt-get -y upgrade
apt-get -y dist-upgrade
apt-get -y autoremove --purge

# Harden SSH: non-default port, no passwords, no root login
sed -i "s/^#\?Port .*/Port ${ssh_port}/" /etc/ssh/sshd_config
sed -i "s/^#\?PasswordAuthentication .*/PasswordAuthentication no/" /etc/ssh/sshd_config
sed -i "s/^#\?PermitRootLogin .*/PermitRootLogin no/" /etc/ssh/sshd_config
sed -i "s/^#\?ChallengeResponseAuthentication .*/ChallengeResponseAuthentication no/" /etc/ssh/sshd_config

# Fail2Ban + unattended upgrades + prereqs
apt-get install -y fail2ban unattended-upgrades ca-certificates curl gnupg lsb-release

# Ensure unattended upgrades is enabled
systemctl enable --now unattended-upgrades

# --- nftables firewall (default drop) ---
nft flush ruleset || true
nft add table inet filter
nft add chain inet filter input   { type filter hook input   priority 0 ; policy drop ; }
nft add chain inet filter forward { type filter hook forward priority 0 ; policy drop ; }
nft add chain inet filter output  { type filter hook output  priority 0 ; policy accept ; }

# Allow loopback and established/related
nft add rule inet filter input iif lo accept
nft add rule inet filter input ct state established,related accept

# Restrict service ports to trusted CIDRs (IPv4)
for ip in ${allowed_cidr_blocks}; do
  nft add rule inet filter input ip saddr "$ip" tcp dport ${ssh_port} accept
  nft add rule inet filter input ip saddr "$ip" tcp dport 443 accept
  nft add rule inet filter input ip saddr "$ip" tcp dport ${portainer_port} accept
done

# Allow HTTP only if explicitly enabled
if [ "${enable_http}" = "true" ]; then
  for ip in ${allowed_cidr_blocks}; do
    nft add rule inet filter input ip saddr "$ip" tcp dport 80 accept
  done
fi

# Persist rules and lock permissions
nft list ruleset > /etc/nftables.conf
chmod 600 /etc/nftables.conf
systemctl enable --now nftables

# Apply SSH changes
systemctl daemon-reexec
systemctl restart ssh

# --- Docker (from official repo) ---
install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg
chmod a+r /etc/apt/keyrings/docker.gpg
echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" \
  | tee /etc/apt/sources.list.d/docker.list > /dev/null

apt-get update
apt-get install -y docker-ce docker-ce-cli containerd.io docker-compose-plugin

systemctl enable --now docker
usermod -aG docker ubuntu

# Optional: pre-pull Portainer CE
docker pull portainer/portainer-ce:latest