#!/bin/bash
set -euo pipefail

# shrink clonezilla images by repacking them with better compression
# v1.0, ed <irc.rizon.net>, MIT-licensed, https://ocv.me/dev/clonezilla-repk.sh

# the compression command / settings to use
compress="zstd -czv --long -T0 -19"

##
## end of configuration section

msg()  { printf '\033[0;36;7m*\033[27m %s\033[0m%s\n' "$*" >&2; }
inf()  { printf '\033[1;92;7m+\033[27m %s\033[0m%s\n' "$*" >&2; }
warn() { printf '\033[0;33;7m!\033[27m %s\033[0m%s\n' "$*" >&2; }
err()  { printf '\033[1;91;7mx\033[27m %s\033[0m%s\n' "$*" >&2; }
die()  { err "$*"; exit 1; }
absreal() { realpath "$1" || readlink -f "$1"; }

td=$(mktemp --tmpdir -d czrepk.XXXXX || mktemp -d -t czrepk.XXXXX || mktemp -d)
trap "rm -rf $td" INT TERM EXIT
pdir=
pdir2=

fin() {
	[ "$pdir" ] || return 0
	# verify and update checkums of recompressed ptcl's in image
	for sname in MD5SUMS SHA1SUMS B2SUMS; do
		[ -e "$pdir/$sname" ] || continue
		cp -np "$pdir/$sname" "$pdir/repk/$sname" || true
		paste $td/1$sname $td/1fn | while read hash _ fn; do
			awk -v fn="$fn" -v ln="$hash  $fn" '/SUMS$/{next} $2!=fn{print;next} $0!=ln{err=1} END{exit err}' "$pdir/repk/$sname" >"$pdir/repk/$sname.2" ||
				die "input file corrupt; image $sname does not contain the hash computed just now: $hash $fn"

			mv "$pdir/repk/$sname"{.2,}
		done
		paste $td/2$sname $td/2fn | while read hash _ fn; do
			printf '%s  %s\n' "$hash" "$fn" >>"$pdir/repk/$sname"
		done
		msg "updated $pdir/$sname"
	done

	szo=$(find "$pdir"      -maxdepth 1 -printf '%s %p\n' | awk '/-ptcl-img\./{v+=$1} END{print v}')
	szn=$(find "$pdir/repk" -maxdepth 1 -printf '%s %p\n' | awk '/-ptcl-img\./{v+=$1} END{print v}')
	mb=$(echo "scale=2;($szo-$szn)/(1024*1024)" | bc)
	perc=$(awk </dev/null "END{print(int(10000-10000*($szn/$szo))/100)}")

	rm -v -- "$pdir/"*.*-ptcl-img.*
	mv "$pdir/repk/"* "$pdir/"
	rmdir "$pdir/repk"
	printf '%s\n' "$compress" >"$pdir/compression"
	inf "finished repacking $pdir from $((szo/(1024*1024))) to $((szn/(1024*1024))) MiB, saved $mb MiB, $perc%"
	echo
	pdir=
	rm $td/{1,2}*
}

# default to all images below working directory, unless there are script args;
# assume script args are relative/absolute paths and recursively do those instead
msg "recursively scanning for clonezilla images to recompress..."
srcdirs=(.)
[ "$#" -gt 0 ] && srcdirs=("$@")
for d in "${srcdirs[@]}"; do
	d="$(absreal "$d")"
	msg "scanning $d"
	find "$d" -type f -iname '*.*-ptcl-img.*' |
	sed -r 's/(.*-ptcl-img\.[^\.]+)\.[a-z]{2,3}$/\1/'
done | sort | uniq > $td/ptcls
# old clonezillas have a different splitsuffix, so change if necessary: [a-z]{2,3}

nptcls=$(wc -l <$td/ptcls)
nimgs=$(sed -r 's/[^/]+$//' <$td/ptcls | uniq | wc -l)
msg "found $nimgs images with $nptcls disk partitions in total"

nptcl=0
nimg=0
while IFS= read -r ipath; do
	nptcl=$((nptcl+1))
	dir="${ipath%/*}"
	srcname="${ipath##*/}"
	basename="${srcname%.*}"
	sc=${ipath##*.}
	dec=$(case $sc in
		gz) echo zcat;;
		bz2) echo bzcat;;
		lzo) echo lzop -dc;;
		lzma) echo lzcat;;
		xz) echo xzcat;;
		lzip) echo lzip -dc;;
		lrz) echo lrzip -d;;
		lz4) echo lz4cat;;
		zst) echo zstdcat;;
		uncomp) echo cat;;
	esac)
	[ "$dec" ] || die "ERROR unknown compression [$sc]"

	[ -e "$dir/clonezilla-img" ] || {
		warn "skipping $ipath because there is no 'clonezilla-img' file next to it"
		continue
	}

	# if this ptcl is from another folder/image, finish up any previously processed image befoer continuing
	[ "$dir" = "$pdir" ] || {
		fin
		nimg=$((nimg+1))
		grep -qF "$compress" "$dir/compression" 2>/dev/null && {
			[ "$pdir2" = "$dir" ] &&
				nimg=$((nimg-1)) ||
				msg "skipping image $nimg/$nimgs (already compressed): $dir"

			pdir2="$dir"
			continue
		}
		inf "started processing image $nimg/$nimgs, $dir"
		mkdir "$dir/repk" || die "ERROR found partially recompressed image; manually delete the repk folder to allow a retry"
	}

	sz=$(awk -v "v=$basename" '/^Running:/{o=0} index($0,v){o=1} o&&/^Space in use:/{print $4 " " $5}' "$dir/clonezilla-img")
	msg "recompressing part $nptcl/$nptcls (image $nimg/$nimgs), $srcname (uncompressed size: $sz)"
	#sz=$(for f in "$ipath"*; do cat <"$f"; done | $dec | pv -i0.5 | wc -c | awk '{print$1}')

	# recompress the ptcl while verifying the checksums of the input files and hashing the decompressed output
	for f in "$ipath"*; do
		printf '%s\n' "${f##*/}" >>$td/1fn
		cat "$f" | tee >(md5sum >>$td/1MD5SUMS) >(sha1sum >>$td/1SHA1SUMS) >(b2sum >>$td/1B2SUMS)
	done | $dec | tee >(b2sum >$td/1dec) |
	$compress |
	split -a2 -b1024M - "$dir/repk/$basename.zst." || die "ERROR recompression pipeline failed"
	# clonezilla-webdav reads each file into ram; splitsize selected accordingly

	# verify image using partclone.chkimage if available and applicable (it can't do dd)
	verifycmd="partclone.chkimg -s - -L /dev/null --no_block_detail"
	command -v partclone.chkimg >/dev/null || verifycmd=pv
	printf '%s\n' "$basename" | grep -qE '\.dd-ptcl-img$' && verifycmd=pv

	# verify recompression; grab checksums of each file and their decompressed output
	for f in "$dir/repk/$basename.zst."*; do
		printf '%s\n' "${f##*/}" >>$td/2fn
		tee <"$f" >(md5sum >>$td/2MD5SUMS) >(sha1sum >>$td/2SHA1SUMS) >(b2sum >>$td/2B2SUMS)
	done | zstdcat | tee >(b2sum >$td/2dec) | $verifycmd >/dev/null

	# verify the decompressed data is identical
	cmp $td/{1,2}dec || die "ERROR recompressed data corrupted"
	pdir="$dir"
done < <(cat $td/ptcls)
fin
inf "all done with $nimgs images, $nptcls disk partitions"
