#!/bin/bash
# Fix root SSH access on all provisioned machines.
# Tries root, lab, michal users to find one that works,
# then ensures root has the SSH key and PermitRootLogin is enabled.
set -euo pipefail

SSH_KEY="ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQDMJ3FkUGbG174eoO5RjZd2eNV680FM5pgp0AgpW/QwlJExK3qxMk0DJSr4ICmzGUx4yujAXcrqU1otcOMPzzFzwc5heWpSmlNHU3TIW6NHEt0sF9ZTAbGLw2zSw3si5UouqFkCcENA40mePFJqY+Q9R8N1uvLgu4m/do+Zrn/mk5Ewc1V7OCRE5Acrnaec4T7LTB0BuVXcjPUfAmZ0q5fI+bKPR1q2Kc3+IeGhVkBuZ9OJVeXXhnpedm0uEbLeriK/jUYKYw/1QhsNDM8Tyty+UIGr9QVnWwzCMHB+wuQcDYC9mPGTqg0fYwX8Mp8xMi1PPxdsh1G7bj/cpWMAF43KswWORF2ul8ICGbaE1zEgIYXO790SuBjpBHhaC6Iegqi58hmCuP+a9893q/EU9HyrWTJHCZXC5E4kP1MsM57KrhEpszM6I3sW9f9zMTPd5QsCXFi4si4OMwX4kYNVu3fQGQPpseDPlTTSrT6uUdqj4Irm0c1m9cYTmK0vYgsM3ss= michal@fedora"

SSH_OPTS="-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR -o ConnectTimeout=5"
USERS_TO_TRY=(root lab michal)

# Machines: hostname ip
MACHINES=(
    "labmaster 192.168.8.11"
    "worker0-k8s0 192.168.8.23"
    "worker1-k8s0 192.168.8.13"
    "worker2-k8s0 192.168.8.25"
    "spark-2935 192.168.8.12"
)

BOLD="\033[1m"
GREEN="\033[0;32m"
RED="\033[0;31m"
DIM="\033[2m"
RESET="\033[0m"

# Script to run on each machine (via sudo if needed)
read -r -d '' FIX_SCRIPT << 'FIXEOF' || true
#!/bin/bash
set -e
KEY="$1"

# 1. Ensure root .ssh dir exists
mkdir -p /root/.ssh
chmod 700 /root/.ssh
touch /root/.ssh/authorized_keys
chmod 600 /root/.ssh/authorized_keys

# 2. Add key if not present
if ! grep -qF "$KEY" /root/.ssh/authorized_keys 2>/dev/null; then
    echo "$KEY" >> /root/.ssh/authorized_keys
    echo "KEY_ADDED"
else
    echo "KEY_EXISTS"
fi

# 3. Fix sshd_config for root login with keys
SSHD_CONF="/etc/ssh/sshd_config"
CHANGED=0

# Ensure PermitRootLogin allows key auth
CURRENT=$(grep -E "^PermitRootLogin" "$SSHD_CONF" 2>/dev/null | tail -1 || true)
if [ "$CURRENT" = "PermitRootLogin prohibit-password" ] || [ "$CURRENT" = "PermitRootLogin without-password" ]; then
    echo "SSHD_OK"
elif [ "$CURRENT" = "PermitRootLogin yes" ]; then
    echo "SSHD_OK"
else
    # Remove any existing PermitRootLogin lines
    sed -i '/^#*PermitRootLogin/d' "$SSHD_CONF"
    echo "PermitRootLogin prohibit-password" >> "$SSHD_CONF"
    CHANGED=1
    echo "SSHD_FIXED"
fi

# Ensure PubkeyAuthentication is enabled
if grep -qE "^PubkeyAuthentication no" "$SSHD_CONF" 2>/dev/null; then
    sed -i 's/^PubkeyAuthentication no/PubkeyAuthentication yes/' "$SSHD_CONF"
    CHANGED=1
    echo "PUBKEY_FIXED"
else
    echo "PUBKEY_OK"
fi

# Restart sshd if changed
if [ "$CHANGED" -eq 1 ]; then
    systemctl restart sshd 2>/dev/null || systemctl restart ssh 2>/dev/null || true
    echo "SSHD_RESTARTED"
fi

# 4. Verify root can be reached
echo "DONE"
FIXEOF

echo ""
echo -e "${BOLD}Fixing root SSH access on all machines...${RESET}"
echo ""

for entry in "${MACHINES[@]}"; do
    read -r hostname ip <<< "$entry"
    printf "  %-24s ${DIM}(%s)${RESET}  " "$hostname" "$ip"

    # Try each user until one works
    WORKING_USER=""
    for user in "${USERS_TO_TRY[@]}"; do
        if ssh $SSH_OPTS "$user@$ip" "true" 2>/dev/null; then
            WORKING_USER="$user"
            break
        fi
    done

    if [ -z "$WORKING_USER" ]; then
        echo -e "${RED}UNREACHABLE${RESET} (tried: ${USERS_TO_TRY[*]})"
        continue
    fi

    # Run fix script (with sudo if not root)
    if [ "$WORKING_USER" = "root" ]; then
        RESULT=$(ssh $SSH_OPTS "root@$ip" "bash -s -- '$SSH_KEY'" <<< "$FIX_SCRIPT" 2>&1)
    else
        RESULT=$(ssh $SSH_OPTS "$WORKING_USER@$ip" "sudo bash -s -- '$SSH_KEY'" <<< "$FIX_SCRIPT" 2>&1)
    fi

    # Parse result
    DETAILS=""
    if echo "$RESULT" | grep -q "KEY_ADDED"; then DETAILS="key added"; fi
    if echo "$RESULT" | grep -q "KEY_EXISTS"; then DETAILS="key ok"; fi
    if echo "$RESULT" | grep -q "SSHD_FIXED"; then DETAILS="$DETAILS, sshd fixed"; fi
    if echo "$RESULT" | grep -q "SSHD_OK"; then DETAILS="$DETAILS, sshd ok"; fi
    if echo "$RESULT" | grep -q "SSHD_RESTARTED"; then DETAILS="$DETAILS, restarted"; fi

    # Verify root works now
    if ssh $SSH_OPTS "root@$ip" "true" 2>/dev/null; then
        echo -e "${GREEN}OK${RESET} ${DIM}(via $WORKING_USER: $DETAILS)${RESET}"
    else
        echo -e "${RED}PARTIAL${RESET} ${DIM}(via $WORKING_USER: $DETAILS -- root still blocked)${RESET}"
    fi
done

echo ""
echo -e "${BOLD}Done.${RESET} Verify: labctl provision recheck --user root"
echo ""