#!/bin/bash
set -e

# create clonezilla usb (bios/uefi) with some tweaks / better defaults
# ed <irc.rizon.net>, MIT-licensed, https://ocv.me/dev/clonezilla-usb.sh
#
# testing with qemu:
# truncate -s 256M test.hdd && echo ,,0c | sfdisk -q --label dos -w always test.hdd && mkfs.ext4 -E offset=1048576 -d testdata test.hdd
# truncate -s 1G cz.usb && ./clonezilla-usb.sh clonezilla-live-3.1.0-22-amd64.zip cz.usb y
# qemu-system-x86_64 -enable-kvm -cpu host -bios /usr/share/OVMF/OVMF_CODE.fd -device virtio-blk-pci,drive=cz,bootindex=1 -drive id=cz,if=none,format=raw,file=cz.usb -drive format=raw,discard=on,detect-zeroes=unmap,file=test.hdd -m 2048


# repack initramfs and squashfs for 60% faster boot
optimize=1

# support non-uefi machines (set 1 to enable; default is uefi-only)
bios=


# --- end of config section --- begin program ---
zip="$1"
dev="$2"
[ "$bios" = 1 ] || bios=
[ "$optimize" = 1 ] || optimize=
die() { printf '%s\n' "$*"; exit 1; }

traps=
addtrap() {
	traps="$1; $traps"
	trap "$traps" INT TERM EXIT
}

[ $(id -u) -eq 0 ] || die need root

err=
printf '%s\n' "$zip" | grep -qiE "zip$" || {
	echo "need arg 1: clonezilla zip file, for example clonezilla-live-3.1.0-22-amd64.zip"
	err=1
}
[ "$dev" ] && [ -e "$dev" ] && [ ! -b "$dev" ] && {
	echo "target is not a blockdevice; fixing with losetup"
	rmmod loop || true
	modprobe loop max_part=63
	dev=$(losetup -f --show "$dev")
	addtrap "losetup -d $dev"
}
printf '%s\n' "$dev" | grep -qE "^/dev/(sd[a-z]|loop[0-9])" || {
	echo "need arg 2: blockdevice to write, for example /dev/sdg"
	err=1
}
[ $err ] && die "optional arg 3: 'y' to skip yes/no checks"

dr=$dev
is_part=
echo $dev | grep -qE 'loop[0-9]p[0-9]|sd[a-z][0-9]' && is_part=1
[ $is_part ] && dr=$(echo $dev | sed -r 's/(sd[a-z]|loop[0-9]+).*/\1/')
if echo $dev | grep -q loop; then
	dp1=${dr}p1; dp2=${dr}p2
else
	dp1=${dr}1; dp2=${dr}2
fi
[ $is_part ] && [ $dev != $dp2 ] && die "must be whole device or the second partition (got $dev, want $dp2)"

echo
lsblk $dev -o NAME,SIZE,TRAN,TYPE,MODEL,MOUNTPOINTS && echo
blkid $dev* && echo
echo "will destroy all data on the above target; confirm by pressing Y:"
# ask for confirmation unless 3rd argument is "y"
[ "$3" = y ] || {
	read -u1 -n1 -r; echo
	[ "$REPLY" = y ] || exit 1
}
set -x
umount $dev* || true
[ $is_part ] || {
	head -c $((1024*1024*2)) /dev/zero > $dr
	sz=$(lsblk $dr -nbo SIZE | awk 'NR==1{print int($1/(1024*1024))}')
	[ $bios ] && bootable=',*' || bootable=
	printf '%s\n' ,+$((sz-513))M,0c ",,ef$bootable" | sfdisk -q --label dos -w always $dr
	mkfs.vfat -S512 --mbr=n -nCZ_IMG $dp1 >/dev/null
}
mkfs.vfat -S512 --mbr=n -nCZ_BOOT $dp2 >/dev/null
td=$(mktemp -d)
td2=$(mktemp -d)
addtrap "[ -e $td2/m ] && umount $td2/m; umount $td; rmdir $td; rm -rf $td2"
mount $dp2 $td
excl=(
	.disk/\*
	live/freedos.img
	live/filesystem.squashfs
	EFI/boot/boot\*
	boot/grub/efi.img
	boot/grub/i386-efi/\*
	boot/grub/unicode.pf2
	boot/grub/\*.png
)
[ $bios ] && excl+=(
	syslinux/\*.png
	utils/win\*
) || excl+=(
	syslinux/\*
	utils/\*
)
unzip -q "$zip" -d $td -x "${excl[@]}"

is_ubuntu=$(awk 'NR==1&&/-[0-9]{8}-/{print 1}' $td/Clonezilla-Live-Version)

mv $td/EFI/boot/{grub,boot}x64.efi  # secureboot-shim is busted; use grub directly

[ $bios ] && {
	patch $td/utils/linux/makeboot.sh <<'EOF'
@@ -90 +90 @@
-    read continue_choice 
+    read -n1 continue_choice 
@@ -100 +100 @@
-            echo "$msg_you_have_to_enter_yes_or_no!"
+            echo "$msg_you_have_to_enter_yes_or_no!"; continue_choice=
@@ -191 +191 @@
-  local RE2='i2o/hd[a-z]+[0-9]+'
+  local RE2='loop[0-9]+p[0-9]'
EOF
	[ "$3" = y ] && a=-b || a=
	(cd $td/utils/linux; ./makeboot.sh $a $dp2)
}

# repack initramfs for 2sec faster boot (27->25),
# repack squashfs for 8sec faster boot (25->17)
unzip -q "$zip" -d $td2 live/initrd.img live/filesystem.squashfs
if [ $optimize ]; then
	[ $is_ubuntu ] &&
		cp $td2/live/initrd.img $td/live/initrd.img ||
		xz -d < $td2/live/initrd.img | zstd --long -T0 -19 > $td/live/initrd.img &

	mkdir $td2/m
	mount $td2/live/filesystem.squashfs $td2/m
	(cd $td2/m && find -type d | cut -c3- | grep -E '/modules/.*/kernel.*/(wireless|bluetooth|mac80211|infiniband|sound|amdgpu|media|netfilter|ipv6)$|/espeak-ng') > $td2/e
	mksquashfs $td2/m $td/live/filesystem.squashfs -ef $td2/e -b 1024k -comp zstd -Xcompression-level 19 -exit-on-error
	umount $td2/m
	wait
	rm -rf $td2
else
	mv $td2/live/* $td/live/
fi

# set kernel args
tf1=$(mktemp)
tf2=$(mktemp)
cat $td/boot/grub/grub.cfg >$tf1
cat $td/syslinux/syslinux.cfg >$tf2 || true
for kv in \
	locales=en_US.UTF-8 \
	keyboard-layouts=NONE \
	ocs_prerun=/run/live/medium/opt.sh \
	ocs_live_type=device-image \
	ocsroot_src=local_dev \
;do
	IFS== read k v <<<$kv
	for tf in $tf1 $tf2; do
		sed -ri "s& $k=[^ ]* ?& &; s&( boot=live .*)&\1 $k=$v&" $tf
	done
done
# for tf in $tf1 $tf2; do sed -ri 's/ quiet loglevel=0 / /' $tf; done  # debug
sed -ri 's`^play .*`play 1920 330 1 0 2 330 1`; s`cyan/blue`cyan/black`; s`^(set timeout=").*`\12"`' $tf1
sed -ri 's`\x07``; /MENU (BACKGROUND|COLOR)/d; s`^(timeout )[0-9]+`\120`' $tf2
cat $tf1 > $td/boot/grub/grub.cfg
cat $tf2 > $td/syslinux/syslinux.cfg || true
rm $tf1 $tf2

# hack the stuff that can't be done the right way
# ======================================================================
cat >$td/opt.sh <<'OPT_EOF'
#!/bin/bash

# assume USB devices are already present, don't wait
patch /usr/sbin/ocs-prep-repo <<'EOF'
@@ -42,2 +42,3 @@
           --msgbox "$msg_prompt_for_insert_USB_dev_if_necessary" 0 0
+   elif true; then true
    else
@@ -72,2 +73,3 @@
    # Question about fsck the repo partition.
+   echo fsck-y >$ANS_TMP ||
    $DIA --backtitle "$msg_nchc_free_software_labs" --title \
EOF

# prefer CZ_IMG (first flashdrive partition) for images
patch /usr/share/drbl/sbin/ocs-functions <<'EOF'
@@ -4924,4 +4924,5 @@
     while [ "$ASK_INPUT" -eq 1 ]; do
+      local defsrc=$(df /run/live/medium/ --output=source | awk '{d=$1} END{sub(/.*\//,"",d);sub(/2/,"1",d);print d}')
       $DIA --backtitle "$msg_nchc_free_software_labs" --title "$msg_nchc_clonezilla | $msg_mode: $ocs_mode_prompt" \
-        --$selection "$msg_choose_dev_to_what" 0 0 $height $HARDDEVS \
+        --$selection "$msg_choose_dev_to_what" 0 0 $height --default-item "$defsrc" $HARDDEVS \
       2> $TMP
EOF

# mount top of filesystem for images, don't browse
patch /usr/sbin/ocs-live-bind-mount <<'EOF'
@@ -137,2 +137 @@
-  BrowseCurrentDirectory "$pre_ocsroot_path" $sel_dir
-  real_ocsroot="$(cat $sel_dir)"
+  real_ocsroot="$pre_ocsroot_path"
EOF

# skip pointless pauses
patch /usr/sbin/ocs-sr <<'EOF'
@@ -882,4 +882,3 @@
     if [ "$ocs_batch_mode" != "on" ]; then
-      echo -n "$msg_press_enter_to_continue "
-      read
+      true
     fi
EOF
sed -ri 's/(countdown )[0-9]+$/\13/' /usr/sbin/ocs-live-final-action 

# safe because ttyS0 is disabled, only TTY1 is active
patch /usr/share/drbl/sbin/ocs-functions <<'EOF'
@@ -11338,2 +11337,3 @@
   trap "[ -f "$TMP" ] && rm -f $TMP" HUP INT QUIT TERM EXIT
+  echo Start_Clonezilla >$TMP ||
   $DIA --backtitle "$msg_nchc_free_software_labs" --title  \
EOF

# better defaults
patch /usr/share/drbl/sbin/drbl-functions <<'EOF'
@@ -1482,3 +1482,3 @@
   "$msg_mode: $ocs_mode_prompt" --menu "$msg_choose_post_mode_after_clone_1:" \
-  0 0 0 $DIA_ESC \
+  0 0 0 --default-item '-p poweroff ' $DIA_ESC \
   "-p choose "    "$msg_ocs_param_p_choose_1" \
@@ -1811,3 +1811,3 @@
       "$msg_clonezilla_advanced_extra_param | $msg_mode: $ocs_mode_prompt" --menu "$msg_choose_if_fsck_the_source_part" \
-      0 0 0 $DIA_ESC \
+      0 0 0 --default-item '-fsck-y ' $DIA_ESC \
       " "     "$msg_skip_check_source_fs" \
@@ -1877,3 +1877,3 @@
       "$msg_clonezilla_advanced_extra_param | $msg_mode: $ocs_mode_prompt" --menu "$msg_choose_one_compression_param_to_save" \
-      0 0 0 $DIA_ESC \
+      0 0 0 --default-item "${ocs_z9p_option_1:--z1p}" $DIA_ESC \
       "-z1p"     "$msg_ocs_param_z1p" \
@@ -1904,3 +1904,3 @@
         "$msg_clonezilla_advanced_extra_param | $msg_mode: $ocs_mode_prompt" --menu "$msg_choose_one_compression_param_to_save" \
-        0 0 0 $DIA_ESC \
+        0 0 0 --default-item "${ocs_z9p_option_1:--z1p}" $DIA_ESC \
         "-z1p"              "$msg_ocs_param_z1p" \
@@ -2223,3 +2223,3 @@
       "$msg_clonezilla_advanced_extra_param | $msg_mode: $ocs_mode_prompt" --menu "$msg_choose_one_compression_param_to_save" \
-      0 0 0 $DIA_ESC \
+      0 0 0 --default-item "${ocs_z9p_option_1:--z1p}" $DIA_ESC \
       "-z1p"           "$msg_ocs_param_z1p" \
@@ -2248,3 +2248,3 @@
         "$msg_clonezilla_advanced_extra_param | $msg_mode: $ocs_mode_prompt" --menu "$msg_choose_one_compression_param_to_save" \
-        0 0 0 $DIA_ESC \
+        0 0 0 --default-item "${ocs_z9p_option_1:--z1p}" $DIA_ESC \
         "-z1p"              "$msg_ocs_param_z1p" \
@@ -2263,3 +2263,3 @@
       "$msg_clonezilla_advanced_extra_param | $msg_mode: $ocs_mode_prompt" --menu "$msg_choose_one_compression_param_to_save" \
-      0 0 0 $DIA_ESC \
+      0 0 0 --default-item "${ocs_z9_option_1:--z1}" $DIA_ESC \
       "-z1"            "$msg_ocs_param_z1" \
@@ -2282,3 +2282,3 @@
         "$msg_clonezilla_advanced_extra_param | $msg_mode: $ocs_mode_prompt" --menu "$msg_choose_one_compression_param_to_save" \
-        0 0 0 $DIA_ESC \
+        0 0 0 --default-item "${ocs_z9_option_1:--z1}" $DIA_ESC \
         "-z1"              "$msg_ocs_param_z1" \
@@ -2318,3 +2318,3 @@
     "$msg_clonezilla_advanced_extra_param | $msg_mode: $ocs_mode_prompt" --menu "$msg_choose_if_fsck_the_source_part" \
-    0 0 0 $DIA_ESC \
+    0 0 0 --default-item '-fsck-y ' $DIA_ESC \
     "-sfsck "  "$msg_skip_check_source_fs" \
EOF

# don't ask about encryption
echo encrypt_ocs_img=-senc >> /etc/ocs/ocs-live.conf 

# confirm_before_clone is funky
patch /usr/share/drbl/sbin/ocs-functions <<'EOF'
@@ -3710 +3710 @@
-  if [ "$confirm_before_clone" = "yes" ]; then
+  if false; then
@@ -3802 +3802 @@
-  if [ "$confirm_before_clone" = "yes" ] && [ "$rest_conf" = "yes" ]; then
+  if false; then
EOF

# ignore drive-names shuffling between reboots (the OS should use UUIDs anyways)
patch /usr/share/drbl/sbin/ocs-functions <<'EOF'
@@ -8588,5 +8588,5 @@
   echo "$msg_restored_OS_change_fstab_grub_conf"
   [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
-  if [ "$ocs_batch_mode" != "on" ]; then
+  if false; then
     echo -n "$msg_press_enter_to_continue..."
     read
EOF

# should be redundant with the huge drbl-functions patch below but just in case either of them stop working
patch /usr/sbin/ocs-sr << 'EOF'
@@ -40 +40 @@
-gen_sha1sum="no"
+gen_sha1sum="yes"
@@ -42 +42 @@
-check_sha1sum="no"
+check_sha1sum="yes"
@@ -50 +50 @@
-save_restore_error_log="no"
+save_restore_error_log="yes"
@@ -61 +61 @@
-nogui="off"   # Default to turn on TUI
+nogui="on"   # Default to turn on TUI
EOF

patch /usr/share/drbl/sbin/drbl-functions <<'EOF'
@@ -1627 +1627 @@
-      "-nogui"   "$msg_ocs_param_nogui" off \
+      "-nogui"   "$msg_ocs_param_nogui" on \
@@ -1644 +1644 @@
-      "-icds"    "$msg_ocs_param_icds" off \
+      "-icds"    "$msg_ocs_param_icds" on \
@@ -1649 +1649 @@
-      "-cs"      "$msg_ocs_param_cs" off \
+      "-cs"      "$msg_ocs_param_cs" on \
@@ -1658 +1658 @@
-      "-srel"    "$msg_ocs_param_srel" off \
+      "-srel"    "$msg_ocs_param_srel" on \
@@ -1667 +1667 @@
-        echo "-g auto" "-e1 auto" "-e2" "-r" "-x" "-j2" >> $OCS_PARAM_TMP
+        echo "-g auto" "-e1 auto" "-e2" "-nogui" "-r" "-x" "-j2" "-icds" "-cs" "-srel" >> $OCS_PARAM_TMP
@@ -1673 +1673 @@
-        echo "-e1 auto" "-e2" "-t" "-r" "-x" >> $OCS_PARAM_TMP
+        echo "-e1 auto" "-e2" "-nogui" "-t" "-r" "-x" "-cs" "-srel" >> $OCS_PARAM_TMP
@@ -1786 +1786 @@
-      "-nogui "    "$msg_ocs_param_nogui" off \
+      "-nogui "    "$msg_ocs_param_nogui" on \
@@ -1795 +1795 @@
-      "-gs "       "$msg_ocs_param_gs" off \
+      "-gs "       "$msg_ocs_param_gs" on \
@@ -1806 +1806 @@
-      echo "-j2 " >> $OCS_PARAM_TMP
+      echo "-j2 -nogui -gs " >> $OCS_PARAM_TMP
@@ -2010 +2010 @@
-    "-nogui"   "$msg_ocs_param_nogui" off \
+    "-nogui"   "$msg_ocs_param_nogui" on \
@@ -2015 +2015 @@
-    "-c"       "$msg_ocs_param_c" on \
+    "-c"       "$msg_ocs_param_c" off \
@@ -2027 +2027 @@
-    "-icds"    "$msg_ocs_param_icds" off \
+    "-icds"    "$msg_ocs_param_icds" on \
@@ -2032 +2032 @@
-    "-cs"      "$msg_ocs_param_cs" off \
+    "-cs"      "$msg_ocs_param_cs" on \
@@ -2038 +2038 @@
-    "-srel"    "$msg_ocs_param_srel" off \
+    "-srel"    "$msg_ocs_param_srel" on \
@@ -2053 +2053 @@
-      echo "-g auto" "-e1 auto" "-e2" "-r" "-j2" "$opt_confirm_" >> $OCS_PARAM_TMP
+      echo "-g auto" "-e1 auto" "-e2" "-nogui" "-r" "-j2" "-icds" "-cs" "-srel" "$opt_confirm_" >> $OCS_PARAM_TMP
@@ -2059 +2059 @@
-      echo "-e1 auto" "-e2" "-t" "-r" "$opt_confirm_" >> $OCS_PARAM_TMP
+      echo "-e1 auto" "-e2" "-nogui" "-t" "-r" "-icds" "-cs" "-srel" "$opt_confirm_" >> $OCS_PARAM_TMP
@@ -2156 +2156 @@
-    "-c "       "$msg_ocs_param_c" on \
+    "-c "       "$msg_ocs_param_c" off \
@@ -2158 +2158 @@
-    "-nogui "   "$msg_ocs_param_nogui" off \
+    "-nogui "   "$msg_ocs_param_nogui" on \
@@ -2165 +2165 @@
-    "-gs "      "$msg_ocs_param_gs" off \
+    "-gs "      "$msg_ocs_param_gs" on \
@@ -2174 +2174 @@
-    echo "-c " "-j2 " >> $OCS_PARAM_TMP
+    echo "-c " "-j2 -nogui -gs " >> $OCS_PARAM_TMP
EOF
OPT_EOF
# ======================================================================

cp "$0" $td/

cat <<'EOF'

almost done -- writing to disk...
wait until prompt appears before unplugging the USB flashdrive:

EOF
