This script automates the setup of a WSL2 Ubuntu environment with ZSH, Oh My Zsh, Azure CLI, and other helpful tools and configurations.

prep.sh

#!/usr/bin/env bash
set -euo pipefail

info(){ printf "[INFO] %s\n" "$*"; }
warn(){ printf "[WARN] %s\n" "$*"; }
err(){ printf "[ERROR] %s\n" "$*" >&2; }

command_exists(){ command -v "$1" >/dev/null 2>&1; }

# Check if script was run with sudo and provide guidance
check_sudo(){
  if [ "$(id -u)" -eq 0 ]; then
    # Script is running as root
    if [ -n "${SUDO_USER:-}" ]; then
      # Running with sudo
      real_user="$SUDO_USER"
      real_home=$(getent passwd "$SUDO_USER" | cut -d: -f6)
      info "Running as root via sudo. Real user: $real_user"
      # Set environment variables for the rest of the script
      export REAL_USER="$SUDO_USER"
      export REAL_HOME="$real_home"
      export SUDO_DETECTED=true
    else
      # Logged in as root directly
      warn "Running as root directly, not via sudo. This may cause permission issues."
      export REAL_USER="root"
      export REAL_HOME="/root"
      export SUDO_DETECTED=true
    fi
  else
    # Not running as root
    info "Running as regular user: $(whoami)"
    export REAL_USER="$USER"
    export REAL_HOME="$HOME"
    export SUDO_DETECTED=false
  fi
}

is_wsl(){ grep -qi microsoft /proc/version 2>/dev/null || false; }

detect_distro(){
  if [ -r /etc/os-release ]; then
    . /etc/os-release
    echo "$ID"
  else
    unameOut=$(uname -s)
    case "${unameOut}" in
      Darwin) echo "macos" ;;
      *) echo "unknown" ;;
    esac
  fi
}

install_zsh(){
  if command_exists zsh; then
    info "zsh already installed: $(command -v zsh)"
    return
  fi
  
  if is_wsl; then
    info "Detected WSL environment - optimizing for Ubuntu on WSL2"
  fi

  if command_exists apt-get; then
    info "Installing zsh via apt"
    sudo apt-get update -y
    sudo apt-get install -y zsh curl git
  else
    warn "apt-get not found. Cannot install zsh automatically on this system. Please install zsh manually."
    return
  fi

  if command_exists zsh; then
    info "zsh installed successfully: $(command -v zsh)"
  else
    err "zsh installation failed"
  fi
}

install_oh_my_zsh(){
  # Always install Oh My Zsh as the real user, not as root
  if [ "${SUDO_DETECTED:-false}" = true ] && [ "$REAL_USER" != "root" ]; then
    info "Installing Oh My Zsh for user $REAL_USER"
    # Check if Oh My Zsh is already installed for the real user
    if [ -d "$REAL_HOME/.oh-my-zsh" ]; then
      info "Oh My Zsh already installed for $REAL_USER"
      return
    fi
    
    # Run the installer as the real user with sudo -u
    export RUNZSH=no
    export CHSH=no
    
    if command_exists curl; then
      sudo -u "$REAL_USER" sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)"
    elif command_exists wget; then
      sudo -u "$REAL_USER" sh -c "$(wget -qO- https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)"
    else
      err "curl or wget required to install Oh My Zsh"
      return
    fi
    
    info "Oh My Zsh install finished for $REAL_USER"
  else
    # Original code for non-sudo runs
    if [ -d "$HOME/.oh-my-zsh" ]; then
      info "Oh My Zsh already installed"
      return
    fi
  
    info "Installing Oh My Zsh (unattended: will not change default shell or start zsh)"
    export RUNZSH=no
    export CHSH=no
  
    if command_exists curl; then
      sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)"
    elif command_exists wget; then
      sh -c "$(wget -qO- https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)"
    else
      err "curl or wget required to install Oh My Zsh"
      return
    fi
  
    info "Oh My Zsh install finished"
  fi
}

install_azure_cli(){
  if command_exists az; then
    info "Azure CLI already installed: $(command -v az)"
    return
  fi

  if command_exists apt-get; then
    info "Installing Azure CLI (Debian/Ubuntu)"
    curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash
  else
    warn "apt-get not found. Falling back to pip user install for azure-cli (not recommended)"
    if command_exists python3 && command_exists pip3; then
      python3 -m pip install --user azure-cli
      info "You may need to add ~/.local/bin to your PATH to use 'az'"
    else
      err "Cannot install Azure CLI: neither apt-get nor python3/pip3 found"
      return
    fi
  fi

  if command_exists az; then
    info "Azure CLI installed successfully: $(command -v az)"
  else
    warn "Azure CLI may require PATH update or manual intervention"
  fi
}

install_wslu(){
  if command_exists wslview; then
    info "WSLU already installed: $(dirname "$(command -v wslview)")"
    return
  fi

  if ! is_wsl; then
    warn "Not running in WSL environment. Skipping WSLU installation."
    return
  fi

  if command_exists apt-get; then
    info "Installing WSLU (WSL Utilities)"
    # Update sources list if needed (Ubuntu-specific)
    if [ -f /etc/apt/sources.list.d/wsl-translinux.list ] || [ -f /etc/apt/sources.list.d/wslu.list ]; then
      info "WSLU repository already configured"
    else
      info "Adding WSLU repository"
      # For Ubuntu 20.04 and later, wslu is in the default repositories
      # But add specific repo to ensure latest version
      if command_exists add-apt-repository; then
        sudo add-apt-repository -y ppa:wslutilities/wslu
      else
        warn "add-apt-repository not found, attempting to install software-properties-common"
        sudo apt-get update -y
        sudo apt-get install -y software-properties-common
        sudo add-apt-repository -y ppa:wslutilities/wslu
      fi
    fi

    sudo apt-get update -y
    sudo apt-get install -y wslu
  else
    warn "apt-get not found. Cannot install WSLU automatically."
    return
  fi

  if command_exists wslview; then
    info "WSLU installed successfully: $(dirname "$(command -v wslview)")"
  else
    err "WSLU installation failed"
  fi
}

setup_browser_alias(){
  if ! is_wsl; then
    return
  fi

  # Handle both sudo and non-sudo cases for browser alias
  if [ "${SUDO_DETECTED:-false}" = true ] && [ "$REAL_USER" != "root" ]; then
    # When running as sudo, modify the real user's files
    for rc_file in "$REAL_HOME/.bashrc" "$REAL_HOME/.zshrc"; do
      if [ -f "$rc_file" ] && ! grep -q "alias open=" "$rc_file"; then
        # Use sudo -u to write to the file as the real user
        echo "" | sudo -u "$REAL_USER" tee -a "$rc_file" > /dev/null
        echo "# Simple browser opener for WSL" | sudo -u "$REAL_USER" tee -a "$rc_file" > /dev/null
        echo "alias open='wslview'" | sudo -u "$REAL_USER" tee -a "$rc_file" > /dev/null
        info "Added 'open' alias to $rc_file"
      fi
    done
  else
    # Original code for non-sudo runs
    for rc_file in "$HOME/.bashrc" "$HOME/.zshrc"; do
      if [ -f "$rc_file" ] && ! grep -q "alias open=" "$rc_file"; then
        echo "" >> "$rc_file"
        echo "# Simple browser opener for WSL" >> "$rc_file"
        echo "alias open='wslview'" >> "$rc_file"
        info "Added 'open' alias to $rc_file"
      fi
    done
  fi
}

fix_wsl_interop(){
  if ! is_wsl; then
    info "Not running in WSL, skipping WSLInterop fix"
    return
  fi

  info "Checking WSLInterop configuration..."
  
  # Check if WSLInterop is already working
  if [ -e /proc/sys/fs/binfmt_misc/WSLInterop ]; then
    info "WSLInterop is already configured correctly"
    return
  fi
  
  info "WSLInterop is missing. Fixing automatically..."
  
  # Step 1: Install binfmt-support
  info "Installing binfmt-support..."
  if command_exists apt-get; then
    sudo apt-get update
    sudo apt-get install -y binfmt-support
  else
    err "apt-get not found. Cannot install binfmt-support."
    return
  fi
  
  # Step 2: Create WSLInterop.conf
  info "Creating WSLInterop configuration..."
  sudo mkdir -p /usr/lib/binfmt.d
  sudo sh -c 'printf ":WSLInterop:M::MZ::/init:PF" > /usr/lib/binfmt.d/WSLInterop.conf'
  
  # Step 3: Restart systemd-binfmt.service
  info "Restarting binfmt service..."
  if command_exists systemctl; then
    sudo systemctl restart systemd-binfmt.service
    info "systemd-binfmt.service restarted"
  else
    # For systems without systemctl
    if [ -x /etc/init.d/systemd-binfmt ]; then
      sudo /etc/init.d/systemd-binfmt restart
    else
      warn "systemctl not found and no init script available. Manual restart needed."
    fi
  fi
  
  # Check if fix was successful
  if [ -e /proc/sys/fs/binfmt_misc/WSLInterop ]; then
    info "WSLInterop has been successfully fixed!"
  else
    warn "WSLInterop is still not available. WSL restart required."
    info "Please follow these steps to restart WSL after the script completes:"
    info "1. Save any work and close all WSL terminals"
    info "2. Open Windows PowerShell or Command Prompt"
    info "3. Run the command: wsl --shutdown"
    info "4. Start WSL again"
  fi
}

print_post_steps(){
  printf "\nSummary:\n"
  command_exists zsh && printf " - zsh: installed at %s\n" "$(command -v zsh)" || printf " - zsh: NOT installed\n"
  [ -d "$HOME/.oh-my-zsh" ] && printf " - Oh My Zsh: installed at %s\n" "$HOME/.oh-my-zsh" || printf " - Oh My Zsh: NOT installed\n"
  command_exists az && printf " - Azure CLI (az): installed at %s\n" "$(command -v az)" || printf " - Azure CLI: NOT installed\n"
  command_exists wslview && printf " - WSLU: installed at %s\n" "$(dirname "$(command -v wslview)")" || printf " - WSLU: NOT installed\n"
  [ -e /proc/sys/fs/binfmt_misc/WSLInterop ] && printf " - WSLInterop: AVAILABLE\n" || printf " - WSLInterop: NOT available\n"
  
  # Show shell information and instructions
  current_shell=${SHELL-}
  if command_exists zsh; then
    zsh_path=$(command -v zsh)
    if [ "$current_shell" = "$zsh_path" ]; then
      printf "\nYour default shell is already set to zsh (%s).\n" "$zsh_path"
    else
      printf "\nTo change your default shell to zsh:\n"
      printf "1. Run the following command:\n"
      printf "   chsh -s %s\n" "$zsh_path"
      printf "2. Enter your password when prompted\n"
      printf "3. Log out and log back in, or restart your terminal\n\n"
    fi
  fi
}

main(){
  # Check if running as sudo/root
  check_sudo
  
  # New order as requested:
  # 1. Install Azure CLI
  # 2. Install WSLU
  # 3. Fix WSL Interop
  # 4. Setup Browser Alias
  # 5. Install ZSH
  # 6. Install Oh My ZSH
  
  info "Starting environment setup in WSL..."
  install_azure_cli
  install_wslu
  fix_wsl_interop
  setup_browser_alias
  install_zsh
  install_oh_my_zsh
  print_post_steps
  
  # Set up alias for current session only
  if command_exists wslview; then
    # Just define the alias for the current session
    alias open="wslview"
    info "The 'open' command is now available in this session."
    
    if [ "${SUDO_DETECTED:-false}" = true ]; then
      info "NOTE: You ran this script with sudo. The 'open' alias is only active for the root user in this session."
      info "To use the alias as $REAL_USER, start a new terminal or run: source ~/.zshrc"
    else
      info "For persistent changes, restart your terminal or run: source ~/.zshrc"
    fi
  fi
  
  # If run with sudo, provide instructions for the user
  if [ "${SUDO_DETECTED:-false}" = true ]; then
    info "Script was run with sudo. All installations should be complete."
    info "Remember that some user-specific settings were applied to user: $REAL_USER"
  fi
  
  # Print shell changing instructions at the very end so they're clearly visible
  if command_exists zsh; then
    zsh_path=$(command -v zsh)
    target_user="${REAL_USER:-$USER}"
    
    # Get the real user's shell when running with sudo
    if [ "${SUDO_DETECTED:-false}" = true ] && [ "$REAL_USER" != "root" ]; then
      # Get the real user's shell from passwd
      if command_exists getent; then
        current_shell=$(getent passwd "$REAL_USER" | cut -d: -f7 || true)
      else
        # Fallback - assume bash if we can't determine
        current_shell="/bin/bash"
      fi
    else
      # Regular user - use SHELL environment variable
      current_shell=${SHELL-}
    fi
    
    # Only show instructions if user's shell is not already zsh
    if [ "$current_shell" != "$zsh_path" ]; then
      printf "\n\033[1;32m==== HOW TO CHANGE YOUR DEFAULT SHELL TO ZSH ====\033[0m\n"
      printf "For user: \033[1m%s\033[0m (current shell: %s)\n\n" "$target_user" "$current_shell"
      
      printf "\033[1m1. Run this command:\033[0m\n"
      printf "   chsh -s %s\n\n" "$zsh_path"
      printf "\033[1m2. Enter your password when prompted\033[0m\n\n"
      printf "\033[1m3. Log out and log back in, or restart your terminal\033[0m\n\n"
    fi
  fi
}

main "$@"