#!/usr/bin/env bash
set -euo pipefail

ROOTFS="${ROOTFS_DIR:-/opt/nvidia/Linux_for_Tegra/rootfs}"

if [[ ! -d "$ROOTFS" ]]; then
    echo "ERROR: Rootfs directory not found: $ROOTFS" >&2
    exit 1
fi

echo "=== Hardening rootfs: $ROOTFS ==="

echo "[1/5] Disabling SSH..."
for unit in sshd.service ssh.service; do
    sudo ln -sf /dev/null "$ROOTFS/etc/systemd/system/$unit" 2>/dev/null || true
done
sudo rm -f "$ROOTFS/etc/ssh/sshd_config"

echo "[2/5] Masking getty and serial console services..."
for unit in "getty@.service" "serial-getty@.service"; do
    sudo ln -sf /dev/null "$ROOTFS/etc/systemd/system/$unit"
done

echo "[3/5] Disabling serial console in bootloader config..."
EXTLINUX="$ROOTFS/boot/extlinux/extlinux.conf"
if [[ -f "$EXTLINUX" ]]; then
    sudo sed -i 's/console=ttyTCU0[^ ]*//' "$EXTLINUX"
    sudo sed -i 's/console=ttyS0[^ ]*//' "$EXTLINUX"
    sudo sed -i 's/  */ /g' "$EXTLINUX"
fi

echo "[4/5] Applying sysctl hardening..."
sudo tee "$ROOTFS/etc/sysctl.d/99-azaion-hardening.conf" > /dev/null <<'EOF'
kernel.yama.ptrace_scope = 3
kernel.core_pattern = |/bin/false
kernel.kptr_restrict = 2
kernel.dmesg_restrict = 1
net.ipv4.conf.all.rp_filter = 1
net.ipv4.conf.default.rp_filter = 1
net.ipv4.conf.all.accept_redirects = 0
net.ipv4.conf.default.accept_redirects = 0
net.ipv4.conf.all.send_redirects = 0
net.ipv4.conf.default.send_redirects = 0
EOF

echo "[5/5] Locking root account..."
if [[ -f "$ROOTFS/etc/shadow" ]]; then
    sudo sed -i 's|^root:[^:]*:|root:!:|' "$ROOTFS/etc/shadow"
fi

echo ""
echo "Hardening complete."