RPI_OS_INSTALL_IN_SD_TO_M2.sh

#!/bin/bash

set -e

PACKAGES="raspberrypi-kernel-headers zfs-dkms zfsutils-linux curl cryptsetup lvm2 lsof grub-efi-arm64"
ZFS_MODULE="zfs"
MOUNT_POINT="/mnt/newroot"
IMAGE_URL="https://downloads.raspberrypi.org/raspios_lite_arm64_latest"
IMAGE_FILE="raspios_lite_arm64_latest.img"
COMPRESSED_FILE="${IMAGE_FILE}.xz"

error_exit() {
    echo "Error: $1" >&2
    exit 1
}

check_root() {
    [ "$(id -u)" = "0" ] || error_exit "This script must be run as root."
}

install_packages() {
    apt update && apt install -y $PACKAGES || error_exit "Failed to install necessary packages"
}

prepare_device() {
    local device=$1
    
    # Unmount all partitions
    for partition in $(lsblk -nlo NAME $device | tail -n +2); do
        umount /dev/$partition 2>/dev/null || true
    done
    
    # Disable swap
    swapoff -a
    
    # Deactivate LVM volume groups
    vgchange -an
    
    # Export ZFS pools
    zpool export -a
    
    # Clear the device
    wipefs -af $device
    sgdisk --zap-all $device
    
    # Inform the kernel of partition table changes
    partprobe $device
    
    # Wait for a moment to ensure changes are processed
    sleep 5
}

create_zfs_pool() {
    local device=$1
    local efi_partition="${device}p1"
    local zfs_partition="${device}p2"

    sgdisk -n 1:1M:+512M -t 1:EF00 $device
    sgdisk -n 2:0:0 -t 2:BF01 $device

    modprobe $ZFS_MODULE || error_exit "Failed to load ZFS module"
    zpool create -f -o ashift=12 rpool $zfs_partition || error_exit "ZFS pool creation failed"
    zfs create rpool/root || error_exit "Root filesystem creation failed"
    zfs create rpool/root/boot || error_exit "Boot filesystem creation failed"

    mkfs.vfat -F 32 $efi_partition
}

download_and_extract_image() {
    if [ ! -f "$IMAGE_FILE" ]; then
        if [ ! -f "$COMPRESSED_FILE" ]; then
            echo "Downloading Raspberry Pi OS image..."
            wget -q --show-progress "$IMAGE_URL" -O "$COMPRESSED_FILE" || error_exit "Image download failed"
        fi
        echo "Extracting image..."
        unxz "$COMPRESSED_FILE" || error_exit "Image extraction failed"
    fi
    [ -f "$IMAGE_FILE" ] || error_exit "Image file not found after download/extraction"
}

copy_image_to_zfs() {
    echo "Copying image to ZFS pool..."
    dd if="$IMAGE_FILE" of=/rpool/root/boot/raspios.img bs=4M status=progress || error_exit "Image copy failed"
}

setup_boot() {
    local efi_partition="/dev/nvme0n1p1"
    local efi_mountpoint="/boot/efi"

    if [ ! -d "$efi_mountpoint" ]; then
        echo "Creating $efi_mountpoint directory..."
        mkdir -p "$efi_mountpoint"
    fi

    echo "Mounting EFI system partition..."
    mount -t vfat "$efi_partition" "$efi_mountpoint" || error_exit "Failed to mount EFI system partition"

    echo "root=/dev/zfs" >> /boot/config.txt || error_exit "Boot settings modification failed"
    grub-install --target=arm64-efi --efi-directory="$efi_mountpoint" --bootloader-id=GRUB || error_exit "GRUB installation failed"
    sed -i 's/GRUB_CMDLINE_LINUX_DEFAULT=""/GRUB_CMDLINE_LINUX_DEFAULT="root=ZFS=rpool\/root"/' /etc/default/grub
    update-grub || error_exit "GRUB update failed"

    echo "Unmounting EFI system partition..."
    umount "$efi_mountpoint"
}

main() {
    check_root
    install_packages

    lsmod | grep $ZFS_MODULE || error_exit "ZFS module loading failed"

    echo "Available disk list:"
    lsblk
    read -p "Enter the M.2 SSD device name (e.g., nvme0n1): " M2_DEVICE
    M2_DEVICE="/dev/${M2_DEVICE}"

    [ -e "$M2_DEVICE" ] || error_exit "Specified device does not exist: $M2_DEVICE"

    echo "WARNING: All data on ${M2_DEVICE} will be erased. This operation is irreversible."
    read -p "Are you sure you want to continue? (y/N): " CONFIRM
    [[ $CONFIRM =~ ^[Yy]$ ]] || error_exit "Operation cancelled by user."

    prepare_device $M2_DEVICE
    create_zfs_pool $M2_DEVICE
    download_and_extract_image
    copy_image_to_zfs
    setup_boot

    echo "Installation completed. Please reboot the system to boot from the M.2 SSD."
    read -p "Do you want to reboot now? (y/N): " REBOOT
    [[ $REBOOT =~ ^[Yy]$ ]] && reboot
}

main

Warning: Use ZFS partition