This is a Bash script designed to set up a secure SSH connection to a Colab notebook using ngrok

cloab-ssh.sh

#!/bin/bash

# Function to print error messages
error() {
    echo "ERROR: $1" >&2
    exit 1
}

# Function to print usage
usage() {
    echo "Usage: $0 <ngrok_authtoken> [ssh_public_key_path]"
    echo "Example: $0 your_ngrok_token ~/.ssh/id_rsa.pub"
    exit 1
}

# Check for required argument
if [ -z "$1" ]; then
    usage
fi

NGROK_TOKEN="$1"
SSH_KEY_PATH="${2:-}"

# Install required packages
echo "Installing required packages..."
apt-get update -qq && apt-get install -qq -o=Dpkg::Use-Pty=0 openssh-server pwgen > /dev/null || error "Failed to install packages"

# Download and setup ngrok
echo "Setting up ngrok..."
if [ ! -f "ngrok" ]; then
    wget -nc -q https://bin.equinox.io/c/bNyj1mQVY4c/ngrok-v3-stable-linux-amd64.tgz || error "Failed to download ngrok"
    tar -xvzf ngrok-v3-stable-linux-amd64.tgz || error "Failed to unzip ngrok"
    rm -f ngrok-v3-stable-linux-amd64.tgz  # Clean up the downloaded file
fi
./ngrok authtoken "$NGROK_TOKEN" || error "Failed to set ngrok authtoken"

# Configure SSH server
echo "Configuring SSH server..."
mkdir -p /var/run/sshd
mkdir -p /root/.ssh
chmod 700 /root/.ssh

# Update SSH configuration to enforce key-only authentication
cat > /etc/ssh/sshd_config << EOF
PermitRootLogin yes
PasswordAuthentication no
PubkeyAuthentication yes
AuthorizedKeysFile .ssh/authorized_keys
ChallengeResponseAuthentication no
UsePAM yes
LogLevel DEBUG3
KexAlgorithms +diffie-hellman-group1-sha1,diffie-hellman-group14-sha1
HostKeyAlgorithms +ssh-rsa
PubkeyAcceptedKeyTypes +ssh-rsa
EOF

# Add debug output for SSH keys
echo "Checking SSH key configuration..."
cat /root/.ssh/authorized_keys
echo "SSH key permissions:"
ls -la /root/.ssh/

# Restart SSH service to apply logging changes
service ssh restart

# Setup environment
echo "LD_LIBRARY_PATH=/usr/lib64-nvidia" >> /root/.bashrc
echo "export LD_LIBRARY_PATH" >> /root/.bashrc

# Handle SSH key
if [ -n "$SSH_KEY_PATH" ]; then
    # If SSH key path is provided, use it
    if [ -f "$SSH_KEY_PATH" ]; then
        cat "$SSH_KEY_PATH" >> /root/.ssh/authorized_keys
        chmod 600 /root/.ssh/authorized_keys
    else
        error "SSH public key file not found at $SSH_KEY_PATH"
    fi
else
    # If no SSH key is provided, prompt for manual input
    # Temporarily redirect debug output
    exec 3>&2
    exec 2>/tmp/ssh_debug.log
    
    clear  # Clear the screen
    echo "Please run 'ssh-keygen' on your local machine if you haven't already"
    echo "Waiting for your public key input..."
    echo "Enter the contents of your public key file (~/.ssh/id_rsa.pub):"
    read -p "> " ssh_key
    
    # Restore debug output
    exec 2>&3
    
    echo "$ssh_key" >> /root/.ssh/authorized_keys
    chmod 600 /root/.ssh/authorized_keys
fi

# Start ngrok
echo "Starting ngrok tunnel..."
./ngrok tcp 22 > /dev/null &

# Wait for ngrok to start
sleep 5

# Get connection details
echo "Getting SSH connection details..."
connection_info=$(curl -s http://localhost:4040/api/tunnels | python3 -c "import sys, json; print('ssh root@' + json.load(sys.stdin)['tunnels'][0]['public_url'].replace('tcp://', ''))")

echo "==============================================="
echo "Connection established!"
echo "Use the following command to connect:"
echo "$connection_info"
echo "Debug information:"
echo "- Checking SSH service status..."
service ssh status
echo "- Checking authorized_keys permissions..."
ls -la /root/.ssh/authorized_keys
echo "- Checking SSH configuration..."
sshd -T
echo "==============================================="