#! /bin/bash

# (c) 2023 by Thomas Lange, lange@debian.org
#
# mk-data-partition  ISO | <USB device>
# extend an ISO file by size and create a data partition
# for USB stick the data partition will take the whole size of the USB stick

mntdir=/media/data

add_p3() {

    # $1: filename or device name
    # $2: partition name
    local type

    if [ -b $2 ]; then
        echo "Partition 3 already exists on $1"
        exit 4
    fi

    if [ X$fmt = Xexfat ]; then
        type=7
    else
        type=83
    fi
    # add 3th partion to the image or device

    { echo -e "n\np\n3\n\n\n"; sleep 2; echo -e "v\nt\n3\n$type\nw\n"; } | fdisk -w never $1 >/dev/null 2>&1
    if [ "$?" -ne 0 ]; then
        echo "Cannot add 3th partition to $1. Error code $?"
        exit 3
    fi
    udevadm settle --timeout=10
}

make_fs() {

    if [ -b $1 ]; then
        if [ X$fmt = Xexfat ]; then
            mkfs.exfat -L MY-DATA $1
        else
            mkfs.ext4 -q -L MY-DATA -J size=5 -E lazy_itable_init $1
        fi
        echo "Data partition MY-DATA created"
    fi
}

mount_data() {

    mkdir $mntdir
    mount $1 $mntdir
}

umount_data() {

    umount $mntdir
    rmdir $mntdir

}
copy_data() {

    cp -a $* $mntdir
    echo "Copied files into the ISO."
    du --exclude=lost+found -sh $mntdir/*
}


usage() {

    cat <<EOF
Usage: mk-data-partition [OPTION] IMAGE [dir...]

Create additional partition for storing data.

   -F              Format data partition as exFAT. Default is to use ext4
   -s              Set size of data partition. Default 300M, ignored for USB devices
   -c              Copy a list of directories to the data partition
   IMAGE           Can be an ISO file or the device of an USB stick

The command adds a third partition to an ISO or an USB stick containing a bootable ISO.
This partition contains an ext4 or exFAT file system with label MY-DATA that can be
mounted read-write. An ISO file will be extended by the given size, on an USB stick the
partition will take the whole remaining size.

Examples:

    Add a data partition of size 1G to the Debian installer ISO using an ext4 partition
    # mk-data-partition -s 1G debian-12.2.0-amd64-netinst.iso

    Create the data partition using an exFAT file system on USB named /dev/sdb.
    First copy (or dd) the ISO onto the USB stick. Then you can add the data partition
    to the USB stick.
    # mk-data-partition -F /dev/sdb

    Create the data partition and copy directories A and B to it
    # mk-data-partition -c debian-12.2.0-amd64-netinst.iso A B

    Copyright (C) 2023 by Thomas Lange

EOF
    exit $1
}

while getopts Fhs:c opt ; do
    case "$opt" in
        h) usage 0;;
        c) copy=1 ;;
        F) fmt=exfat ;;
        s) size=$OPTARG ;;
        *) usage 1;;
    esac
done
shift $((OPTIND - 1))

filename=$1
shift
if [ -z "$filename" ]; then
    echo "ERROR: No filename or device supplied"
    echo
    usage 1
fi
filename=$(readlink -f $filename)

if expr "$size" : '^[0-9]\+$' >/dev/null; then
   echo "Size $size needs a unit like K, M, G, T"
   exit 5
fi

if [ -n "$1" ] && [ X$copy = X ]; then
    echo "If you want to copy directories, add -c"
    exit 3
fi

if [ $(id -u) != "0" ]; then
    echo "Run this program as root."
    exit 2
fi

if [ ! -b $filename ] && [ ! -f $filename ]; then
    echo "$filename is neither a device nor an existing file"
    exit 9
fi

if [ -b $filename ]; then
    # file is a USB device
    if [ -n "$size" ]; then
        echo "Warning: Ignoring the size parameter for USB sticks."
    fi
    part=${filename}3
    add_p3 $filename $part
    make_fs $part
fi

if [ -f $filename ]; then

    # check if the image already has a 3th partition
    set +e
    fdisk -l $filename | egrep -q ^${filename}3
    if [ $? -eq 0 ]; then
        echo "Partition 3 already exists on $filename"
        exit 5
    fi
    size=${size:-300M}
    echo "Extend $filename by $size"
    truncate -s +$size $filename
    if [ $? -ne 0 ]; then
        echo "Cannot extend $filename. Aborting"
        exit 9
    fi
    loop=$(losetup -P -f --show $filename)
    if [ "$?" -ne 0 ]; then
        echo "Cannot create loop device"
        exit 5
    fi
    trap "losetup -d $loop" EXIT
    part=${loop}p3
    add_p3 $loop $part
    make_fs $part
fi

if [ X$copy = X1 ]; then

    mount_data $part
    copy_data $*
    umount_data $part
fi