#!/bin/bash

# Setup passwordless SSH to a remote server
# Usage: ./setup-ssh.sh <ip_address> <username> [admin_user]

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Check arguments
if [ $# -lt 2 ]; then
    echo -e "${YELLOW}Usage: $0 <ip_address> <username> [admin_user]${NC}"
    echo ""
    echo "Arguments:"
    echo "  ip_address  - Target server IP"
    echo "  username    - User to setup SSH for (will be created if doesn't exist)"
    echo "  admin_user  - (Optional) Existing user with sudo access to create new user"
    echo ""
    echo "Examples:"
    echo "  $0 192.168.200.100 ansible          # Setup for existing user"
    echo "  $0 192.168.200.100 ansible root     # Create 'ansible' user via root"
    exit 1
fi

IP_ADDRESS="$1"
USERNAME="$2"
ADMIN_USER="${3:-}"
SSH_KEY="$HOME/.ssh/id_ed25519"

echo -e "${GREEN}Setting up passwordless SSH to ${USERNAME}@${IP_ADDRESS}${NC}"

# Generate SSH key if it doesn't exist
if [ ! -f "$SSH_KEY" ]; then
    echo -e "${YELLOW}SSH key not found. Generating new ed25519 key...${NC}"
    ssh-keygen -t ed25519 -f "$SSH_KEY" -N "" -C "ansible@$(hostname)"
    echo -e "${GREEN}SSH key generated: ${SSH_KEY}${NC}"
else
    echo -e "${GREEN}Using existing SSH key: ${SSH_KEY}${NC}"
fi

# If admin user provided, create target user if it doesn't exist
if [ -n "$ADMIN_USER" ]; then
    echo -e "${YELLOW}Connecting as ${ADMIN_USER} to setup user ${USERNAME}...${NC}"
    echo "You may be prompted for the password for ${ADMIN_USER}@${IP_ADDRESS}"

    ssh -o StrictHostKeyChecking=accept-new "${ADMIN_USER}@${IP_ADDRESS}" bash -s <<EOF
        set -e

        # Create user if doesn't exist
        if ! id "$USERNAME" &>/dev/null; then
            echo "Creating user: $USERNAME"
            sudo useradd -m -s /bin/bash "$USERNAME"
            echo "User $USERNAME created"
        else
            echo "User $USERNAME already exists"
        fi

        # Setup .ssh directory
        sudo mkdir -p /home/$USERNAME/.ssh
        sudo chmod 700 /home/$USERNAME/.ssh
        sudo touch /home/$USERNAME/.ssh/authorized_keys
        sudo chmod 600 /home/$USERNAME/.ssh/authorized_keys
        sudo chown -R $USERNAME:$USERNAME /home/$USERNAME/.ssh

        # Add to sudo group (optional - for Ansible privilege escalation)
        if ! groups $USERNAME | grep -qE '\b(sudo|wheel)\b'; then
            if getent group sudo &>/dev/null; then
                sudo usermod -aG sudo "$USERNAME"
                echo "Added $USERNAME to sudo group"
            elif getent group wheel &>/dev/null; then
                sudo usermod -aG wheel "$USERNAME"
                echo "Added $USERNAME to wheel group"
            fi
        fi

        # Configure passwordless sudo for the user
        echo "$USERNAME ALL=(ALL) NOPASSWD:ALL" | sudo tee /etc/sudoers.d/$USERNAME > /dev/null
        sudo chmod 440 /etc/sudoers.d/$USERNAME
        echo "Configured passwordless sudo for $USERNAME"
EOF

    # Copy SSH key to the new user
    echo -e "${YELLOW}Copying SSH key to ${USERNAME}@${IP_ADDRESS}...${NC}"
    PUBKEY=$(cat "${SSH_KEY}.pub")
    ssh "${ADMIN_USER}@${IP_ADDRESS}" "echo '$PUBKEY' | sudo tee -a /home/$USERNAME/.ssh/authorized_keys > /dev/null && sudo chown $USERNAME:$USERNAME /home/$USERNAME/.ssh/authorized_keys"

else
    # Standard ssh-copy-id for existing user
    echo -e "${YELLOW}Copying public key to remote server...${NC}"
    echo "You may be prompted for the password for ${USERNAME}@${IP_ADDRESS}"
    ssh-copy-id -i "${SSH_KEY}.pub" "${USERNAME}@${IP_ADDRESS}"
fi

# Test the connection
echo -e "${YELLOW}Testing passwordless SSH connection...${NC}"
if ssh -o BatchMode=yes -o ConnectTimeout=5 "${USERNAME}@${IP_ADDRESS}" "echo 'SSH connection successful'" 2>/dev/null; then
    echo -e "${GREEN}✓ Passwordless SSH setup complete!${NC}"
    echo -e "${GREEN}You can now connect with: ssh ${USERNAME}@${IP_ADDRESS}${NC}"
else
    echo -e "${RED}✗ SSH connection test failed. Please check your setup.${NC}"
    exit 1
fi