dev_pod_init.sh

#!/bin/bash

# SSH
if [[ -n "$POD_PUBLIC_KEY" || -n "$POD_ROOT_PASSWORD" ]]; then
    # 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

    # Check if POD_PASSWORD_ACCESS variable is set to "true"
    if [[ "$POD_PASSWORD_ACCESS" == "true" ]]; 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

        # Display a message indicating that user/password SSH access is enabled
        echo "User/password SSH access is enabled."
    fi

    # Set root password if POD_ROOT_PASSWORD is provided
    if [[ -n "$POD_ROOT_PASSWORD" ]]; then
        echo "root:${POD_ROOT_PASSWORD}" | chpasswd
        echo "Root password has been set."
    fi

    # 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 "${POD_PUBLIC_KEY}" "$HOME/.ssh/authorized_keys"; then
        # Append the public key to authorized_keys
        echo "$POD_PUBLIC_KEY" >> "$HOME/.ssh/authorized_keys"
        echo "Public key from env variable added."
    fi

    # Check if POD_PERMIT_ROOT_LOGIN is enabled
    if [[ "$POD_PERMIT_ROOT_LOGIN" == "true" ]]; then
        # 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."
    fi

    if [[ -z "$LEPTON_DEVPOD_PORT_2222" ]]; then
      echo "ListenAddress 127.0.0.1" >> /etc/ssh/sshd_config
    fi

    # turn off PAM to fix sshd login issue
    sed -i 's/UsePAM yes/UsePAM no/' /etc/ssh/sshd_config
    mkdir /run/sshd
    chmod 0755 /run/sshd
    echo "Starting sshD"
    nohup /usr/sbin/sshd -D -p 2222 > /dev/null 2>&1 &
fi


# Jupyter
if [[ -n "$POD_JUPYTER_PASSWORD" || -n "$LEPTON_ENABLE_JUPYTER" ]]; then
    if pgrep jupyter-lab > /dev/null 2>&1; then
      echo "jupyter already started"
      return
    fi

    # 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

    jupyter lab --generate-config

    address="0.0.0.0"
    if [[ -z "$LEPTON_DEVPOD_PORT_18889" ]]; then
      address="127.0.0.1"
    fi

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

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

    # Set root directory if JUPYTERLAB_WORKSPACES_DIR is provided
    if [[ -n "$JUPYTERLAB_WORKSPACES_DIR" ]]; then
        echo "c.ServerApp.root_dir = '${JUPYTERLAB_WORKSPACES_DIR}'" >> ~/.jupyter/jupyter_lab_config.py
        echo "JupyterLab root directory set to: ${JUPYTERLAB_WORKSPACES_DIR}"
    fi    

    jupyter lab --allow-root > /var/log/jupyter.log 2>&1 &
fi

cat <<'EOF' | sudo tee /usr/local/bin/lepton-set-password > /dev/null
#!/bin/bash
set -e

usage() {
  echo "Usage:"
  echo "  $0 [jupyter] [password] [-y]"
  echo "  -y: Skip confirmation (optional)"
  echo "  If no arguments are provided, interactive mode will be used."
}
force_cleanup=false
target=""
password=""

for arg in "$@"; do
  if [[ "$arg" == "-y" ]]; then
    force_cleanup=true
  elif [[ "$arg" == "jupyter" ]]; then
    target="$arg"
  elif [[ -z "$password" ]]; then
    password="$arg"
  fi
done
if [[ $# -eq 0 ]]; then
  read -p "Do you want to set jupyter password? (y/N): " confirm
  if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
    echo "Aborted."
    exit 0
  fi
  target="jupyter"

  echo "Enter new password:"
  read -s password1
  echo "Confirm new password:"
  read -s password2
else
  if [[ -z "$target" || -z "$password" ]]; then
    usage
    exit 1
  fi
  password1="$password"
  password2="$password"
fi

if [[ "$target" == "jupyter" ]]; then
  if ! pgrep -f jupyter-lab > /dev/null; then
    echo "[WARN] Jupyter Lab is NOT running"
    exit 0
  fi
fi

if [[ "$password1" != "$password2" ]]; then
  echo "Passwords do not match."
  exit 1
fi

if ! $force_cleanup; then
  echo "You are about to reset the $target password and restart the service."
  read -p "Do you want to proceed? (y/N): " confirm
  if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
    echo "Aborted."
    exit 0
  fi
fi
if [[ "$target" == "jupyter" ]]; then
  echo "[INFO] Setting Jupyter password..."
  mkdir -p ~/.jupyter
  config_file=~/.jupyter/jupyter_lab_config.py
  if grep -q '^c.ServerApp.token' "$config_file"; then
    sed -i "s|^c.ServerApp.token *=.*|c.ServerApp.token = '${password1}'|" "$config_file"
  else
    echo "c.ServerApp.token = '${password1}'" >> "$config_file"
  fi

  tokenfile=~/.local/share/jupyter/runtime/jupyter_cookie_secret
  if [[ -f "$tokenfile" ]]; then
    echo "[INFO] Removing persistent Jupyter token file: $tokenfile"
    rm -f "$tokenfile"
    echo "[INFO] Token file removed. Existing sessions will be invalidated."
  fi

  echo "[INFO] Restarting Jupyter server..."
  pkill -f jupyter-lab || true
  nohup jupyter-lab --allow-root > /dev/null 2>&1 &

  echo "[DONE] Jupyter password updated and server restarted."
else
  echo "Invalid target: $target"
  usage
  exit 1
fi
EOF

# Make it executable
sudo chmod +x /usr/local/bin/lepton-set-password

echo "Exposing ENV variables"
env | grep -v '^$' | sed 's/=/="/' | sed 's/$/"/' > /etc/environment
echo "set -a; source /etc/environment; set +a; cd \$PWD;" >> /root/.bashrc
echo "Removing the init script"
rm ./dev_pod_init.sh
sleep infinite