yastfli

#!/usr/bin/env bash
#
#   yastfli -  Yet Another Script To Flash LiveOS ISOs
#
#   Copyright (C) 2024, Manuel Fombuena <mfombuena@innovara.co.uk>
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <https://www.gnu.org/licenses/>.
#

export PATH=/sbin:/usr/sbin:$PATH

### HELP STARTS ###
usage() {
  echo "
 SYNTAX

 (sudo) yastfli [options] --iso <path> --target <device>

 OPTIONS

 --directory|-d
   Use a directory overlay instead of an image file.
   It can't be used with fat32. Size option would be ignored.
   Default is false.

 --filesystem|-f [ext4|fat32]
   Filesystem of the data partition.
   Default is ext4.

 --iso|-i
   Path to Live ISO file.

 --label|-l
   Label of the data partition.
   Default is LIVE.

 --size-mb|-s
   Size in MB of the overlay file.
   It can't be more than 4095 for fat32 data partitions.
   Default is 512.

 --target|-t
   Target device e.g. /dev/sda.

 --help|-h
   This help.
 "
  exit 1
}
### HELP ENDS ###

### root CHECK STARTS ###
if [[ $(id -u) -ne 0 ]]; then
   echo "ERROR:  this script must be run as root."
   exit 2
fi
### root CHECK ENDS ###

### VARIABLE INITIALIZATION STARTS ###
datafs='ext4'
dirovl=false
iso=''
label='LIVE'
size=512
tgt=''
green='\033[0;32m'
nc='\033[0m' # No Color
### VARIABLE INITIALIZATION ENDS ###

### ARGUMENT PROCESSING STARTS ###
while true ; do
  case $1 in
    --help|-h)
      usage
      ;;
    --directory|-d)
      dirovl=true
      shift 0
      ;;
    --filesystem|-f)
      datafs="${2}"
      shift
      ;;
    --iso|-i)
      iso="${2}"
      shift
      ;;
    --label|-l)
      label="${2}"
      shift
      ;;
    --size-mb|-s)
      size="${2}"
      shift
      ;;
    --target|-t)
      tgt="${2}"
      shift
      ;;
    --*|-*)
      echo "ERROR:  invalid syntax. Use -h for help."
      exit 2
      ;;
    *)
      break
  esac
  shift
done
### ARGUMENT PROCESSING ENDS ###

### ARGUMENT HANDLING STARTS ###
if [[ -z "${iso}" || -z "${tgt}" ]]; then
  echo "ERROR:  invalid syntax. You need at least -i <iso> -t <device>."
  exit 2
elif ! [[ -f "${iso}" ]]; then
  echo "ERROR: ISO file ${iso} doesn't exist."
  exit 2
elif ! [[ -b "${tgt}" ]]; then
  echo "ERROR:  ${tgt} not present."
  exit 2
elif [[ "${datafs}" != "ext4" && "${datafs}" != "fat32" ]]; then
  echo "ERROR:  invalid data partition filesystem. It can only be ext4 or fat32."
  exit 2
elif [[ ${dirovl} == true && "${datafs}" == "fat32" ]]; then
  echo "ERROR:  fat32 is not compatible with directory overlay. Use file or change to ext4."
  exit 2
fi
if [[ "${datafs}" == "fat32" && ${size} -gt 4095 ]]; then
  echo "INFO:  fat32 can't hold an overlay of ${size}M. Reduced to 4095M."
  size=4095
fi
tgtp1="${tgt}"1
tgtp2="${tgt}"2
### ARGUMENT HANDLING ENDS ###

### DEVICE PREP STARTS ###
umount "${tgt}"?* &> /dev/null
isoloop=$(losetup -Pf "${iso}" --show)
isodata=$(mktemp -d)
mount "${isoloop}" "${isodata}" &> /dev/null
### DEVICE PREP ENDS ###

### EFI PARTITION STARTS ###
echo -e "${green}[1/7]${nc} Formatting EFI partition"
parted -s "${tgt}" mklabel gpt
parted -s "${tgt}" mkpart '"EFI System Partition"' fat32 1MiB 257MiB
parted -s "${tgt}" set 1 esp on
mkfs.fat -n ESP -F32 "${tgtp1}" &> /dev/null
fsck.fat -aVw "${tgtp1}" &> /dev/null
echo -e "${green}[2/7]${nc} Copying data to EFI partition"
espmnt=$(mktemp -d)
mount "${tgtp1}" "${espmnt}"
rsync -a --progress "${isodata}"/EFI "${isodata}"/images "${espmnt}"/
echo "Syncing filesystem writes to disc. Please wait."
sync -f "${espmnt}"
### EFI PARTITION ENDS ###

### DATA PARTITION STARTS ###
echo -e "${green}[3/7]${nc} Formatting LIVE partition"
parted -s "${tgt}" mkpart '"Main Data Partition"' "${datafs}" 257MiB 100%
case "${datafs}" in
  ext4)
    mkfs.ext4 -F -L "${label}" -O ^64bit -j "${tgtp2}" &> /dev/null
    tune2fs -c0 -i0 -ouser_xattr,acl "${tgtp2}" &> /dev/null
    e2fsck -yf "${tgtp2}" &> /dev/null
    ;;
  fat32)
    mkfs.fat -n "${label}" -F32 "${tgtp2}" &> /dev/null
    fsck.fat -aVw "${tgtp2}" &> /dev/null
    ;;
  *)
    echo "Invalid LIVE file system. It can only be ext4 or fat32."
    exit 1
    ;;
esac
echo -e "${green}[4/7]${nc} Copying data to LIVE partition"
datamnt=$(mktemp -d)
mount "${tgtp2}" "${datamnt}"
rsync -a --progress --exclude="EFI" --exclude="images" "${isodata}/" "${datamnt}"
echo "Syncing filesystem writes to disc. This may take a while..."
sync -f "${datamnt}"/
### DATA PARTITION ENDS ###

### OVERLAY STARTS ###
tgtlabel=($(lsblk -no LABEL "${tgtp2}" || :))
tgtuuid=($(lsblk -no UUID "${tgtp2}" || :))
tgtfile=overlay-"${tgtlabel}"-"${tgtuuid}"
ovlpath="${datamnt}"/LiveOS/"${tgtfile}"
echo -e "${green}[5/7]${nc} Initializing persistent overlay on ${ovlpath}"
if [[ ${dirovl} == true && "${datafs}" == "ext4" ]]; then
  mkdir -m 0755 --context=system_u:object_r:root_t:s0 \
    $datamnt/LiveOS/overlayfs $datamnt/LiveOS/ovlwork
else
  truncate -s 0 "${ovlpath}"
  fallocate -x -l ${size}M "${ovlpath}"
  echo "Syncing filesystem writes to disc. This could take a while."
  sync -f "${ovlpath}"
  mkfs.ext4 -F -L OVERLAY -O ^64bit -j -E nodiscard "${ovlpath}" &> /dev/null
  tune2fs -c0 -i0 -ouser_xattr,acl "${ovlpath}" &> /dev/null
  e2fsck -yf "${tgtp2}" &> /dev/null
  ovl=$(mktemp -d)
  mount "${ovlpath}" "${ovl}"
  mkdir -m 0755 --context=system_u:object_r:root_t:s0 \
    $ovl/overlayfs $ovl/ovlwork
  umount $ovl
  chmod 0600 "${ovlpath}" &> /dev/null || :
fi
### OVERLAY ENDS ###

### GRUB STARTS ###
echo -e "${green}[6/7]${nc} Adding overlay to ${espmnt}/EFI/BOOT/grub.cfg"
isolabel=($(lsblk -no LABEL "${isoloop}" || :))
sed -i -r "s/root=live:CDLABEL="${isolabel}"/root=live:UUID="${tgtuuid}"/" "${espmnt}"/EFI/BOOT/grub.cfg
sed -i -r "s/set default=\"1\"/set default=\"0\"/" "${espmnt}"/EFI/BOOT/grub.cfg
if [[ ${dirovl} == true && "${datafs}" == "ext4" ]]; then
  sed -i -r "s/rd\.live\.image|liveimg/& rd.live.overlay=UUID="${tgtuuid}":\/LiveOS\/overlayfs rd.live.overlay.overlayfs/" "${espmnt}"/EFI/BOOT/grub.cfg
else
  sed -i -r "s/rd\.live\.image|liveimg/& rd.live.overlay=UUID="${tgtuuid}" rd.live.overlay.overlayfs/" "${espmnt}"/EFI/BOOT/grub.cfg
fi
### GRUB ENDS ###

### CLEAN UP STARTS ###
echo -e "${green}[7/7]${nc} Cleaning up"
umount "${isoloop}"
losetup -d "${isoloop}"
umount "${espmnt}"
umount "${datamnt}"
eject "${tgt}"
echo "Finished. You can unplug your device now."
### CLEAN UP ENDS ###