Configure SSH and Jupyter

#!/bin/bash
############ User Space ##############
# Please fill in your SSH public key here (something like ssh-rsa AAAxxxxx lepton@sampleDomain.com)
SSHPUB="YOUR_SSH_PUBLIC_KEY"
# Please fill in your SSH root password here
SSHKEY="YOUR_SSH_ROOT_PASSWORD"
# Please fill in your Jupyter Lab token here, leave empty to skip installation of Jupyter Lab
JUPKEY=""
######################################

export DEBIAN_FRONTEND=noninteractive
export TZ=Etc/UTC

function InstallSSH {
  # Check if OpenSSH server is already installed
  if ! command -v sshd &> /dev/null; then
      echo "OpenSSH server is not installed. Installing..."

      apt update
      apt install -y openssh-server

      echo "OpenSSH server installation complete."
  else
      echo "OpenSSH server is already installed."
  fi

  # Set root password if SSHKEY is provided
  if [[ -n "$SSHKEY" ]]; then
      # Enable password authentication in SSH configuration
      sed -i '/^#PasswordAuthentication/c\PasswordAuthentication yes' /etc/ssh/sshd_config
      sed -i '/^PasswordAuthentication/c\PasswordAuthentication yes' /etc/ssh/sshd_config
      # Enable root login in SSH configuration
      sed -i '/^#PermitRootLogin/c\PermitRootLogin yes' /etc/ssh/sshd_config
      sed -i '/^PermitRootLogin/c\PermitRootLogin yes' /etc/ssh/sshd_config
      echo "Root login is enabled."

      # Display a message indicating that user/password SSH access is enabled
      echo "User/password SSH access is enabled."
      echo "root:${SSHKEY}" | chpasswd
      echo "Root password has been set."
  fi

  # Check if LEPTON_PUBLIC_KEY variable is set and not empty
  if [[ -n "$SSHPUB" ]]; then
      # Create the .ssh directory and authorized_keys file if they don't exist
      if [ ! -d "$HOME/.ssh" ]; then
          mkdir -p "$HOME/.ssh"
          chmod 0700 "$HOME/.ssh"
          echo "Directory $HOME/.ssh created."
      fi
      if [ ! -f "$HOME/.ssh/authorized_keys" ]; then
          touch "$HOME/.ssh/authorized_keys"
          chmod 0600 "$HOME/.ssh/authorized_keys"
          echo "File $HOME/.ssh/authorized_keys created."
      fi
      # Check if the public key is not already present in authorized_keys
      if ! grep -q "${SSHPUB}" "$HOME/.ssh/authorized_keys"; then
          # Append the public key to authorized_keys
          echo "$SSHPUB" >> "$HOME/.ssh/authorized_keys"
          echo "Public key from env variable added."
      fi
  fi

  # turn off PAM to fix sshd login issue
  sed -i 's/UsePAM yes/UsePAM no/' /etc/ssh/sshd_config

  # set default port to 2222
  sed -i 's/#Port 22/Port 2222/' /etc/ssh/sshd_config

  echo "Exposing ENV variables"
  env | sed 's/=/="/' | sed 's/$/"/' > /etc/environment
  echo "set -a; source /etc/environment; set +a;" >> /root/.bashrc

  mkdir -p /run/sshd
  chmod 0755 /run/sshd

  # Start SSH service only if it's not already running
  if ! service sshd status > /dev/null 2>&1; then
    service ssh start
    echo "sshd service started"
  else
    echo "sshd service is already running"
    # Restart to apply configuration changes
    service ssh restart
    echo "sshd service restarted to apply configuration changes"
  fi
}

function InstallJupyter {
    # Check if jupyter lab is already installed
    if ! command -v jupyter-lab &> /dev/null; then
        echo "jupyter lab is not installed. Installing..."

        apt update
        apt install python3 python3-pip -y
        pip install -U virtualenv --break-system-packages
        pip3 install jupyterlab --break-system-packages

        echo "jupyter lab installation complete."
    else
        echo "jupyter lab is already installed."
    fi

    # Check if jupyter is already running and stop it to reconfigure
    if pgrep jupyter-lab > /dev/null 2>&1; then
        echo "jupyter already running, stopping to apply new configuration..."
        pkill jupyter-lab
        sleep 2
    fi

    jupyter lab --generate-config

    {
     echo "c.ServerApp.ip = '0.0.0.0'"
     echo "c.ServerApp.open_browser = False"
     echo "c.ServerApp.port = 18889"
    } >> ~/.jupyter/jupyter_lab_config.py

    # Set token if JUPKEY is provided
    if [[ -n "$JUPKEY" ]]; then
        echo "c.ServerApp.token = '${JUPKEY}'" >> ~/.jupyter/jupyter_lab_config.py
        echo "Jupyter token has been set."
    fi

    jupyter lab --allow-root > /var/log/jupyter.log 2>&1 &
    echo "JupyterLab started with user configuration."
}

if [[ -n $SSHKEY ||  -n $SSHPUB  ]]; then
  InstallSSH
fi

if [[ -n $JUPKEY ]]; then
  InstallJupyter
fi

python3 -c "import torch; print(f'PyTorch version: {torch.__version__}'); print(f'CUDA available: {torch.cuda.is_available()}')"