#! /bin/bash

# (c) 2023-2026 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

set -o pipefail

mntdir=/media/data

die() {

    local e=$1   # first parameter is the exit code
    shift

    echo "$*" >&2 # print error message
    exit $e
}

add_p3() {

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

    [ -b "$2" ] && die 4 "Partition 3 already exists on $1"

    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 || die 3 "Cannot add 3rd partition to $1. Error code $?"
    udevadm settle --timeout=10
}

add_p4() {

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

    [ -b "$2" ] && die 4 "Partition 4 already exists on $1"

    if [ X$fmt = Xexfat ]; then
        type=0700
    else
        type=8300
    fi

    # add 4th gpt partion to the image or device
    {
        echo -e "x\ne\nm\nn\n\n\n\n8300\nw\nY\n" | gdisk $1
    } >/dev/null 2>&1 || die 3 "Cannot add 4th partition to $1. Error code $?"
    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-2026 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)

expr "$size" : '^[0-9]\+$' >/dev/null && die 5 "Size $size needs a unit like K, M, G, T"

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

[ $(id -u) != "0" ] && die 2 "Run this program as root."

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

type=$(fdisk -l $filename 2>/dev/null| grep -Po '(?<=Disklabel type:\s)...' )
case $type in
    dos)
        p=3
        fname=add_p3
        ;;
    gpt)
        p=4
        fname=add_p4
        ;;
esac

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}$p
    $fname $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}$p && die 5 "Partition $p already exists on $filename"
    size=${size:-300M}
    echo "Extend $filename by $size"
    truncate -s +$size $filename || die 9 "Cannot extend $filename. Aborting"
    if ! loop=$(losetup -P -f --show $filename); then
        die 5 "Cannot create loop device"
    fi
    trap "losetup -d $loop" EXIT
    part=${loop}p$p
    $fname $loop $part
    make_fs $part
fi

if [ X$copy = X1 ]; then

    mount_data $part
    copy_data $*
    umount_data $part
fi