diff --git a/autodl-initializer.sh b/autodl-initializer.sh new file mode 100644 index 0000000..5c20450 --- /dev/null +++ b/autodl-initializer.sh @@ -0,0 +1,493 @@ +#!/bin/bash + +# This script is a initialization helper for container +# instances on AutoDL based on my flavour. It does: +# - Set up locale to avoid garbled text display. +# - Set up APT mirrors to speed up package installaion. +# - Install specified packages, then clean up APT caches. +# - Clean up bash history. +# +# Also, there are some tips for your customization and simplifying server configuring. +# - Add your SSH public keys. (Also can be done in AutoDL portal.) +# - Add your custom hook function in this bash, do anything you like. +# - Write files to your destination path. +# - Install `uv`, a powerful replacement for `pip` + + + +############################################################### +# Custom options +# +# Make your customization here! +############################################################### + +# Set the desired system locale. +# Example: 'en_US.UTF-8', 'zh_CN.UTF-8' +SYSTEM_LOCALE="en_US.UTF-8" + +# Set the desired APT mirror URL. Scheme (http/https) is needed. +# Example: 'http://mirrors.kernel.org' +APT_MIRROR="http://mirrors.bfsu.edu.cn" + +APT_SOURCE_FILE="/etc/apt/sources.list" +APT_SOURCES_LIST_D="/etc/apt/sources.list.d" + +# List of packages to install, with spaces separated names. +# Example: "tmux git curl htop" +PACKAGES_TO_INSTALL="tmux git curl nano htop nvtop" + +# SSH public key for authorization. +# Example: "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAII7hPh6ucD76+bdbf4Tyb0uq0lsDbFYVMLFZqmxc8zjP YourKey" +SSH_PUBLIC_KEY="" + +# uv install +UV_DOWNLOAD_URL="https://github.com/astral-sh/uv/releases/latest/download/uv-x86_64-unknown-linux-gnu.tar.gz" +UV_INSTALL_PATH="${HOME}/.local/bin" + + + + +############################################################### +# Custom function hooks +# +# Here you can add your custom functions to the array, +# which will be executed later one by one. +############################################################### + +# The array to hold the function names. +HOOK_FUNCTIONS=() + +# Function to add a function to the hook list. +# Usage: add_hook_function +add_hook_function() { + if [[ -n "$1" ]]; then + HOOK_FUNCTIONS+=("$1") + fi +} + +# You can define your own functions here and add them. +# Example custom function: +# my_custom_setup() { +# echo "This is my custom setup function." +# # Add your custom commands here. +# echo "Setting up a custom user..." +# useradd myuser +# } +# # Add the function to the hook list. +# add_hook_function my_custom_setup +turbo_autodl_network() { + source /etc/network_turbo +} +add_hook_function turbo_autodl_network + +# Download helper, pass the URL and Destination as parameter. +download_file() { + local url=$1 + local output=$2 + + if command -v curl &> /dev/null; then + curl -sSL -o "$output" "$url" || return 1 + elif command -v wget &> /dev/null; then + wget -qO "$output" "$url" || return 1 + else + echo "Error: Neither curl nor wget found. Install one to proceed." >&2 + return 1 + fi +} + +install_uv() { + local url="${UV_DOWNLOAD_URL}" + local install_path="${UV_INSTALL_PATH}" + local shell_type="" + local shell_profile="" + + if ! mkdir -p "$install_path"; then + echo "Error: failed to create install directory: $install_path" >&2 + return 1 + fi + + # curl -sSL "$url" | tar --strip-components=1 -C "$install_path" -xzf - + if ! temp_file=$(mktemp); then + echo "Error: Failed to create temporary file" >&2 + return 1 + fi + + if ! download_file "$url" "$temp_file"; then + rm -f "$temp_file" + return 1 + fi + + if ! tar --strip-components=1 -C "$install_path" -xzf "$temp_file"; then + echo "Error: Failed to extract archive" >&2 + rm -f "$temp_file" + return 1 + fi + rm -f "$temp_file" + + echo "uv installed to $install_path" + + if [ -n "$BASH_VERSION" ]; then + shell_type="bash" + shell_profile="${HOME}/.bashrc" + elif [ -n "$ZSH_VERSION" ]; then + shell_type="zsh" + shell_profile="${HOME}/.zshrc" + else + shell_type="bash" + shell_profile="${HOME}/.profile" + fi + + # Ensure install path is in PATH + if ! echo "$PATH" | grep -q "$install_path"; then + echo "Adding $install_path to PATH and Shell..." + export PATH="$install_path:$PATH" + # Persist for future shells + if ! grep -q "$install_path" "$shell_profile" 2>/dev/null; then + echo "export PATH=\"$install_path:\$PATH\"" >> "$shell_profile" + fi + fi + + echo "Modifying shell profile for uv..." + echo "eval \"\$(uv generate-shell-completion ${shell_type})\"" >> "$shell_profile" + echo "eval \"\$(uvx --generate-shell-completion ${shell_type})\"" >> "$shell_profile" + + echo "Installation complete. Run 'uv --version' to verify." +} +add_hook_function install_uv + + + + +############################################################### +# Custom files +# +# Specify the file contents and their destination. Here you +# can make some configuration so you won't need to do it later. +############################################################### + +# tmux configuration +cat << EOF > "${HOME}/.tmux.conf" +set -g mouse on +set -g default-terminal "screen-256color" +EOF + +# condarc +cat << EOF >> "${HOME}/.condarc" + +auto_activate_base: false +EOF + + + + +############################################################### +# Script functions +# +# These function are for script itself, you can modify them +# as long as you know what you are doing. +############################################################### + +# Function to execute all functions in the hook list. +execute_hooks() { + if [[ ${#HOOK_FUNCTIONS[@]} -ne 0 ]]; then + echo "--- Executing custom hook functions ---" + for func_name in "${HOOK_FUNCTIONS[@]}"; do + if type "$func_name" &>/dev/null; then + echo "--> Running user hook: $func_name" + "$func_name" + else + echo "Warning: Hook function '$func_name' does not exist." + fi + done + fi +} + +# Function to check for root privileges. +check_root() { + if [[ $EUID -ne 0 ]]; then + echo "This script must be run as root." + exit 1 + fi +} + +# Set the system locale. +set_locale() { + echo "--- Setting system locale to '$SYSTEM_LOCALE' ---" + # Generate the locale if it doesn't exist + if ! locale -a | grep -q "$SYSTEM_LOCALE"; then + echo "Generating locale: $SYSTEM_LOCALE" + locale-gen "$SYSTEM_LOCALE" + fi + # Set the system default locale + update-locale LANG="$SYSTEM_LOCALE" LANGUAGE="$SYSTEM_LOCALE" LC_ALL="$SYSTEM_LOCALE" + # echo "Locale setting complete." +} + + +# Function to check for required variables and distro +_check_apt_prerequisites() { + if [[ -z "$APT_MIRROR" ]]; then + echo "Error: The APT_MIRROR variable is not set. Please set it to your desired mirror, e.g., 'mirrors.kernel.org'." + return 1 + fi + + if [[ ! -f "/etc/os-release" ]]; then + echo "Error: Could not detect OS. This script is intended for Debian/Ubuntu based systems." + return 1 + fi + + . /etc/os-release + if [[ "$ID" != "debian" && "$ID" != "ubuntu" ]]; then + echo "Error: This script is intended for Debian or Ubuntu, but detected '$ID'." + return 1 + fi + + return 0 +} + +# Function to detect DEB-822 or legacy sources.list style +_get_source_style() { + if [[ -d "$APT_SOURCES_LIST_D" ]]; then + # Check if there are any .sources files in the sources.list.d directory + if find "$APT_SOURCES_LIST_D" -name "*.sources" -print -quit | grep -q '.*'; then + echo "deb-822" + return + fi + fi + echo "legacy" +} + +# Function to create the new sources.list file +_create_new_sources_list() { + local distro_id="$1" + local distro_version_codename="$2" + local components="" + local security_url_suffix="" + + case "$distro_id" in + "debian") + components="main contrib non-free" + security_url_suffix="-security" + ;; + "ubuntu") + components="main restricted universe multiverse" + security_url_suffix="" + ;; + *) + echo "Error: Unsupported distribution '$distro_id'." + return 1 + ;; + esac + + local mirror_url="${APT_MIRROR}/${distro_id}/" + local security_mirror_url="${APT_MIRROR}/${distro_id}${security_url_suffix}/" + + echo "Creating new $APT_SOURCE_FILE with mirror: $mirror_url" + + cat << EOF > "$APT_SOURCE_FILE" +# +# This file was automatically generated. +# +deb ${mirror_url} ${distro_version_codename} ${components} +deb ${mirror_url} ${distro_version_codename}-updates ${components} +deb ${mirror_url} ${distro_version_codename}-backports ${components} + +# Security updates +deb ${security_mirror_url} ${distro_version_codename}-security ${components} +EOF + + return 0 +} + +# Function to replace APT source for legacy sources.list style +_replace_legacy_sources() { + local distro_id="$1" + local distro_version_codename="$2" + + echo "Detected legacy sources.list format." + echo "Backing up old sources.list to ${APT_SOURCE_FILE}.bak..." + mv "$APT_SOURCE_FILE" "${APT_SOURCE_FILE}.bak" + + _create_new_sources_list "$distro_id" "$distro_version_codename" +} + +# Function to replace APT source for DEB-822 style +_replace_deb822_sources() { + local distro_id="$1" + local distro_version_codename="$2" + local components="" + local security_url_suffix="" + + case "$distro_id" in + "debian") + components="main contrib non-free" + security_url_suffix="-security" + ;; + "ubuntu") + components="main restricted universe multiverse" + security_url_suffix="" + ;; + *) + echo "Error: Unsupported distribution '$distro_id'." + return 1 + ;; + esac + + local new_source_file="${APT_SOURCES_LIST_D}/${distro_id}.sources" + local mirror_url="${APT_MIRROR}/${distro_id}/" + local security_mirror_url="${APT_MIRROR}/${distro_id}${security_url_suffix}/" + + echo "Detected DEB-822 format." + echo "Backing up old sources file to ${new_source_file}.bak..." + [[ -f "$new_source_file" ]] && mv "$new_source_file" "${new_source_file}.bak" + + echo "Creating new DEB-822 style sources file at: $new_source_file" + + cat << EOF > "$new_source_file" +Types: deb +URIs: ${mirror_url} +Suites: ${distro_version_codename} ${distro_version_codename}-updates ${distro_version_codename}-backports +Components: ${components} +Signed-By: /usr/share/keyrings/${distro_id}-archive-keyring.gpg + +Types: deb +URIs: ${security_mirror_url} +Suites: ${distro_version_codename}-security +Components: ${components} +Signed-By: /usr/share/keyrings/${distro_id}-archive-keyring.gpg +EOF + + return 0 +} + +# Replace APT mirror main function. +# Both legacy sources.list and DEB-822 format are suppported. +replace_apt_mirror() { + echo "--- Setting up APT mirrors ---" + _check_apt_prerequisites || return 1 + + . /etc/os-release + local distro_id="$ID" + local distro_version_codename="$VERSION_CODENAME" + + if [[ -z "$distro_version_codename" ]]; then + echo "Error: Could not determine the distribution codename. This is required to proceed." + return 1 + fi + + local source_style=$(_get_source_style) + + case "$source_style" in + "deb-822") + _replace_deb822_sources "$distro_id" "$distro_version_codename" + ;; + "legacy") + _replace_legacy_sources "$distro_id" "$distro_version_codename" + ;; + *) + echo "Error: Unknown source file style detected." + return 1 + ;; + esac + + if [[ $? -eq 0 ]]; then + echo "APT source successfully updated to '$APT_MIRROR'." + echo "Running 'apt update' to refresh package lists..." + apt update + if [[ $? -eq 0 ]]; then + echo "APT update completed successfully." + else + echo "Warning: 'apt update' failed. Please check your network connection and the mirror URL." + fi + else + echo "APT source replacement failed." + return 1 + fi +} + +# Install packages. +install_packages() { + if [[ -z "$PACKAGES_TO_INSTALL" ]]; then + # echo "No packages specified to install. Skipping." + return 0 + fi + echo "--- Installing packages: $PACKAGES_TO_INSTALL ---" + # Use 'apt-get install -y' for non-interactive installation + apt-get install -y $PACKAGES_TO_INSTALL + echo "--- Package installation complete. ---" +} + +add_ssh_key() { + # Check if the SSH_PUBLIC_KEY variable is non-empty + if [[ -n "$SSH_PUBLIC_KEY" ]]; then + echo "--- Adding SSH_PUBLIC_KEY ---" + + local ssh_dir="${HOME}/.ssh" + local authorized_keys_file="${ssh_dir}/authorized_keys" + + # Check if the .ssh directory exists, if not, create it with correct permissions + if [[ ! -d "$ssh_dir" ]]; then + echo "Creating directory: $ssh_dir" + mkdir -p "$ssh_dir" + chmod 700 "$ssh_dir" + fi + + # Check if the authorized_keys file exists, if not, create it with correct permissions + if [[ ! -f "$authorized_keys_file" ]]; then + echo "Creating file: $authorized_keys_file" + touch "$authorized_keys_file" + chmod 600 "$authorized_keys_file" + fi + + # Use grep to check if the public key already exists in the file. + # The 'grep -q' option suppresses output and exits with a 0 status if a match is found. + if ! grep -q -F "$SSH_PUBLIC_KEY" "$authorized_keys_file"; then + echo "Appending key to $authorized_keys_file" + echo "$SSH_PUBLIC_KEY" >> "$authorized_keys_file" + else + echo "Key already exists in $authorized_keys_file." + fi + fi +} + +# Remove APT caches. +clean_apt_cache() { + echo "--- Removing APT caches ---" + rm -rf /var/cache/apt/archives/* /var/lib/apt/lists/* +} + +# Clean bash history. +clean_bash_history() { + echo "--- Cleaning bash history ---" + history -c + # Overwrite the history file with nothing + > ~/.bash_history +} + + + + +############################################################### +# Main process +# +# The main part of this script. +############################################################### + +# Check for root privileges at the beginning. +check_root + +# Main functions +set_locale +replace_apt_mirror +install_packages +clean_apt_cache +add_ssh_key + +# Execute the functions in the hook list. +execute_hooks + +# Clean the bash history as the final step. +clean_bash_history + +# Exit with success status. +exit 0 \ No newline at end of file