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

log() { echo "[$(date +'%Y-%m-%d %H:%M:%S')] $1"; }

if [ "$EUID" -ne 0 ]; then
    log "ERROR: Please run as root"
    exit 1
fi

# Check Tailscale is installed and connected
if ! command -v tailscale &>/dev/null; then
    log "ERROR: Tailscale is not installed"
    exit 1
fi

TAILSCALE_IP=$(tailscale ip -4 2>/dev/null)
if [ -z "$TAILSCALE_IP" ]; then
    log "ERROR: Could not get Tailscale IP — is Tailscale connected?"
    exit 1
fi
log "Tailscale IP: ${TAILSCALE_IP}"

# Detect SSH service name
if systemctl is-active --quiet ssh 2>/dev/null; then
    SSH_SERVICE="ssh"
elif systemctl is-active --quiet sshd 2>/dev/null; then
    SSH_SERVICE="sshd"
else
    log "ERROR: No running SSH service found (tried ssh, sshd)"
    exit 1
fi

# Use a drop-in file if sshd_config.d exists, otherwise edit sshd_config directly
if [ -d /etc/ssh/sshd_config.d ]; then
    DROPIN="/etc/ssh/sshd_config.d/99-tailscale-only.conf"
    [ -f "$DROPIN" ] && log "WARNING: ${DROPIN} already exists — overwriting"
    cat > "$DROPIN" << EOF
# Bind SSH to Tailscale interface only
# To revert: rm ${DROPIN} && systemctl restart ${SSH_SERVICE}
ListenAddress ${TAILSCALE_IP}
PasswordAuthentication no
PermitEmptyPasswords no
EOF
    log "Written: ${DROPIN}"
else
    BACKUP="/etc/ssh/sshd_config.bak.$(date +%Y%m%d_%H%M%S)"
    cp /etc/ssh/sshd_config "$BACKUP"
    log "Backed up sshd_config to ${BACKUP}"

    # Remove any existing ListenAddress lines and add ours
    sed -i '/^[#]*ListenAddress/d' /etc/ssh/sshd_config
    echo "ListenAddress ${TAILSCALE_IP}" >> /etc/ssh/sshd_config
    log "Updated /etc/ssh/sshd_config"
fi

# Validate config before restarting
if ! sshd -t; then
    log "ERROR: SSH config validation failed — aborting without restart"
    [ -n "${DROPIN:-}" ] && rm -f "$DROPIN"
    exit 1
fi
log "SSH config validated OK"

systemctl restart "$SSH_SERVICE"
log "SSH service restarted"

log ""
log "Done. SSH is now bound to Tailscale only (${TAILSCALE_IP})"
log "Connect with: ssh root@${TAILSCALE_IP}"
if [ -d /etc/ssh/sshd_config.d ]; then
    log "To revert:    rm ${DROPIN} && systemctl restart ${SSH_SERVICE}"
else
    log "To revert:    cp ${BACKUP} /etc/ssh/sshd_config && systemctl restart ${SSH_SERVICE}"
fi