#!/usr/bin/env bash

exec > >(tee -i /var/log/stackscript.log) 2>&1
set -euo pipefail

export DEBIAN_FRONTEND=noninteractive
export NEEDRESTART_MODE=a

# <UDF name="NAME" label="Node name" />
# <UDF name="GROUP" label="Create group (Optional)" />
# <UDF name="SSH_USER" label="Create a non-root user" />
# <UDF name="USER_PASSWORD" label="Non-root user password" />
# <UDF name="SSH_PUBLIC_KEY" label="SSH public key for the non-root user" default="" />
# <UDF name="SSH_PORT" label="Set SSH server port" default="22" />
# <UDF name="TIMEZONE" label="Set timezone" default="UTC" />
# <UDF name="ADD_CLOUDFLARE_IPS" label="Add Cloudflare IPs to UFW" type="boolean" default="false" />

touch ~/.hushlogin

echo "Updating system..."
apt-get update

apt-get install -y sudo openssh-server

echo "Setting hostname to $NAME"
hostnamectl set-hostname "${NAME}" || true

: "${SSH_USER:=ansible}"
: "${USER_PASSWORD:=}"

echo "Creating user $SSH_USER"
if ! id -u "${SSH_USER}" >/dev/null 2>&1; then
    useradd -m -s /bin/bash "${SSH_USER}"
fi

if [ -n "${USER_PASSWORD}" ]; then
    echo "${SSH_USER}:${USER_PASSWORD}" | chpasswd
fi

groupadd -f sudo
usermod -aG sudo "${SSH_USER}"
mkdir -p /etc/sudoers.d
cat > "/etc/sudoers.d/90-${SSH_USER}" <<EOF
${SSH_USER} ALL=(ALL) NOPASSWD:ALL
EOF
chmod 440 "/etc/sudoers.d/90-${SSH_USER}"

USER_HOME=$(getent passwd "${SSH_USER}" | cut -d: -f6)
if [ -z "${USER_HOME}" ]; then
    echo "Unable to resolve home directory for user ${SSH_USER}" >&2
    exit 1
fi

if [ -n "${GROUP}" ]; then
    groupadd -f "${GROUP}"
    usermod -aG "${GROUP}" "${SSH_USER}"
fi

# SSH setup
echo "Configuring SSH..."
mkdir -p "${USER_HOME}"/.ssh

for i in $(seq 1 60); do
    if [ -s /root/.ssh/authorized_keys ]; then
        break
    fi
    sleep 2
done

if [ -s /root/.ssh/authorized_keys ]; then
    cp /root/.ssh/authorized_keys "${USER_HOME}"/.ssh/authorized_keys
else
    if [ -n "${SSH_PUBLIC_KEY:-}" ]; then
        printf '%s\n' "${SSH_PUBLIC_KEY}" > "${USER_HOME}"/.ssh/authorized_keys
    else
        echo "No /root/.ssh/authorized_keys and no SSH_PUBLIC_KEY provided" >&2
        exit 1
    fi
fi

if [ -n "${SSH_PUBLIC_KEY:-}" ]; then
    if ! grep -qF "${SSH_PUBLIC_KEY}" "${USER_HOME}"/.ssh/authorized_keys; then
        printf '%s\n' "${SSH_PUBLIC_KEY}" >> "${USER_HOME}"/.ssh/authorized_keys
    fi
fi

chown -R "${SSH_USER}:${SSH_USER}" "${USER_HOME}"/.ssh
chmod 700 "${USER_HOME}"/.ssh
chmod 600 "${USER_HOME}"/.ssh/authorized_keys

chown "${SSH_USER}:${SSH_USER}" "${USER_HOME}"
chmod 755 "${USER_HOME}"
chmod go-w "${USER_HOME}"

mkdir -p /etc/ssh/sshd_config.d
cat > /etc/ssh/sshd_config.d/99-infra.conf <<EOF
Port ${SSH_PORT}
PermitRootLogin no
PasswordAuthentication no
KbdInteractiveAuthentication no
ChallengeResponseAuthentication no
PubkeyAuthentication yes
UsePAM yes
ClientAliveInterval 180
LoginGraceTime 30
MaxAuthTries 3
MaxSessions 10
MaxStartups 10:30:60
StrictModes yes
AuthorizedKeysFile .ssh/authorized_keys
EOF

systemctl restart ssh
systemctl enable ssh

echo "Installing essentials..."
apt-get -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" upgrade
apt-get -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" install ufw fail2ban htop unzip logrotate curl gnupg python3-pip

# Firewall
echo "Configuring UFW firewall..."
ufw default deny incoming
ufw default allow outgoing
ufw allow "${SSH_PORT}/tcp"
ufw limit "${SSH_PORT}/tcp"

ufw --force enable
ufw logging low

mkdir -p /etc/sysctl.d
cat > /etc/sysctl.d/99-console-quiet.conf <<EOF
kernel.printk = 3 4 1 3
EOF
sysctl --system

# Timezone
echo "Setting timezone to ${TIMEZONE}"
timedatectl set-timezone "${TIMEZONE}"

# Docker
echo "Installing Docker..."
apt-get install -y ca-certificates software-properties-common
install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /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" \
  > /etc/apt/sources.list.d/docker.list
apt-get update
apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
usermod -aG docker "${SSH_USER}"
systemctl enable docker

# Docker Compose (v2 plugin installed above)

# Ansible
echo "Installing Ansible..."
pip3 install ansible

# Fail2ban
echo "Configuring Fail2ban..."
cat > /etc/fail2ban/jail.d/sshd.local <<EOF
[sshd]
enabled = true
port = ${SSH_PORT}
maxretry = 3
bantime = 1h
findtime = 10m
EOF
systemctl enable fail2ban
systemctl start fail2ban

# Logrotate
cat > /etc/logrotate.d/custom <<EOF
/var/log/custom/*.log {
    daily
    missingok
    rotate 14
    compress
    delaycompress
    notifempty
    create 0640 root utmp
    sharedscripts
    postrotate
        systemctl reload rsyslog > /dev/null 2>/dev/null || true
    endscript
}
EOF

# Optional: NTP (Systemd handles this well now)
timedatectl set-ntp true

# Cleanup
echo "Cleaning up..."
history -c
rm -f /root/.bash_history /home/${SSH_USER}/.bash_history || true
unset NAME GROUP SSH_USER USER_PASSWORD SSH_PUBLIC_KEY SSH_PORT TIMEZONE ADD_CLOUDFLARE_IPS

echo "StackScript complete. Server ready."