#!/bin/bash
shopt -s nullglob
PATH=/sbin:$PATH

msg() { printf '%s: %s: %s\n' "$0" "$1" "$2" >&2; }
die() { msg Error "${*:-Exiting on fatal error.}"; exit 1; }
warn() { msg Warning "${*:-Something is wrong.}"; }
notice() { msg Notice "$*"; }
quietly() { "$@" >/dev/null 2>&1 || true; }
loudly() { (set -x; "$@"); }

validate_name()
{
        case "$1" in
        *[^a-zA-Z0-9_]*) false ;;
        *) true ;;
        esac
}

read_config_file()
{
        validate_name "$img" || { warn "invalid name: $img"; return 1; }
        while read line
        do
                line=${line%%#*} # ignore comments
                k=${line%%=*}
                v=${line#*=}
                [ "$k" -a "$k" != "$line" ] || return
                eval "conf_${img}_$k=\$v"
        done < "$img".conf
}

inquire_var() { _inquire_var "$img" "$1"; }
_inquire_var()
{
        local v
        v=conf_${1}_${2}
        v=${!v}
        if [ "$v" ]
        then
                eval "$2=\$v"
        else
                false
        fi
}

require_var() { _require_var "$img" "$1"; }
_require_var()
{
        _inquire_var "$@" || die "Missing required field '$2' for image file '$1'"
}

get_root_hash()
{
        sed -ne 's/^Root hash:[ \t]*//p' "$1"
}

require_exists()
{
        local f
        for f
        do
                [ -f "$f" ] || die "Not a file: $f"
        done
}

build_partition_image()
{
        if inquire_var rebuild
        then
                case "$rebuild" in
                        always|never|default) ;;
                        *) die "invalid value for field 'rebuild': $rebuild" ;;
                esac
        fi

        if [ "$rebuild" = 'always' ] || [ ! -e "$imgfile" -a "$rebuild" != 'never' ]
        then

                if [ -e "$imgfile" ]
                then
                        notice "Image file exists: $imgfile"
                fi

                case "$type" in
                        dm-verity-hashes|dm-verity-data)
                                require_var data_path
                                require_exists "$data_path" "$data_path".verity "$data_path".verity.log
                                root_hash=$(get_root_hash "$data_path".verity.log)
                                [ ${#root_hash} = 64 ]
                        ;;
                        *)
                                require_var allocation
                                tmp=$imgfile~tmp
                                fallocate -l "$allocation" "$tmp"
                        ;;
                esac

                case "$type" in
                        efi-system-partition) mkfs.fat -F 32 -I "$tmp" || die "mkfs.vfat failed" ;;
                        bios-grub)            mkfs.fat -F 32 -I "$tmp" || die "mkfs.vfat failed" ;;
                        samizdat-keys)        mkfs.btrfs -q "$tmp" || die "mkfs.btrfs failed" ;;
                        dm-verity-data)
                                partuuid=${root_hash:0:32}
                                cp -f -T --reflink "$data_path" "$builddir"/"$partuuid"
                                ln -sfT "$partuuid" "$tmp"
                        ;;
                        dm-verity-hashes)
                                partuuid=${root_hash:32:32}
                                cp -f -T --reflink "$data_path".verity "$builddir"/"$partuuid"
                                ln -sfT "$partuuid" "$tmp"
                        ;;
                esac
                mv -T "$tmp" "$imgfile"
                notice "Successfully wrote $imgfile"
        fi
}

iterate_partitions()
{
        local f
        for f in part*.conf
        do
                notice "Processing $f"
                img=${f%.conf}

                read_config_file || warn "Received error return from command: read_config_file $img"
                require_var name

                require_var type
                case "$type" in
                        efi-system-partition|bios-grub|samizdat-*) ;;
                        dm-verity-hashes|dm-verity-data) require_var data_path ;;
                        partition-table) ;;
                        *) die "invalid type: $type" ;;
                esac

                imgfile=$builddir/$img

                "$@" || return
        done
}

set_var()
{
        eval "conf_${img}_${1}=\$${1}"
}

create_dmsetup_map()
{
        part=$builddir/${f%.conf}

        if [ -h "$part" ]
        then
                partuuid=$(readlink -e "$part")
        else
                partuuid=
        fi

        $sudo losetup -L -f "$part"
        dev=$(losetup -j "$part" -O NAME --noheadings)
        devsz=$($sudo blockdev --getsz $dev)

        alignment_error=$((devsz % 2048))

        printf '%d %d linear %s 0\n' $start $devsz "$dev" >> "$map"

        set_var start
        set_var devsz
        start=$((start + devsz))

        if [ $alignment_error -gt 0 ]
        then
                devsz=$((2048 - alignment_error))
                printf '%d %d zero\n' $start $devsz >> "$map"
                start=$((start + devsz))
        fi
}

create_ptable_conf_debug()
{
        set -x
        create_ptable_conf "$@"
        set +x
}

create_ptable_conf()
{
        #inquire_var start && inquire_var devsz || return
        part=$builddir/${f%.conf}
        devsz=$(stat -L -c '%s' "$part") || return
        [ "$start" -a "$devsz" ] || return

        [ "$((devsz % 512))" -eq 0 ]

        case "$type" in
                partition-table) start=$((start + devsz)); return;;
                efi-system-partition) typecode=C12A7328-F81F-11D2-BA4B-00A0C93EC93B ;;
                bios-grub)            typecode=21686148-6449-6E6F-744E-656564454649 ;;
                *)                    typecode=0FC63DAF-8483-4772-8E79-3D69D8477DE4 ;;
        esac

        printf '%d: start=%d, size=%d, type=%s, name="%s"\n' \
                "$i" \
                "$((start / 512))" \
                "$((devsz / 512))" \
                "$typecode" \
                "$name" \
                >> "$ptable.sfdisk"
        let ++i
        start=$((start + devsz))
}

set -e

builddir=_build
mkdir -p "$builddir"

iterate_partitions build_partition_image

if [ "$UID" = 0 ]
then sudo=
else sudo=sudo
fi

if whole=$(losetup -j whole.img -O NAME --noheadings)
then
        quietly $sudo umount *.mnt
        quietly $sudo kpartx -dv "$whole"
        quietly $sudo losetup -D
fi

sfdisk_init()
{
        local DEV LAST LAST_LBA
        DEV=$1
        if [ -b "$DEV" ]
        then
                LAST=$($sudo blockdev --getsize64 "$DEV") || return
        else
                LAST=$(stat -L -c%s "$DEV") || return
        fi
        LAST_LBA=$((LAST / 512 - 34))

        cat <<END
label: gpt
unit: sectors
first-lba: 34
last-lba: ${LAST_LBA}

END
}

pee_on_table()
{
        local dev="$1"
        ptable=$builddir/ptable.bin
        i=1
        start=0

        sfdisk_init "$dev" > "$ptable".sfdisk

        iterate_partitions create_ptable_conf_debug

        cat "$ptable".sfdisk >&2

        set -x
        $sudo sfdisk "$dev" < "$ptable".sfdisk || return
        $sudo sfdisk --dump "$dev" || return
        $sudo sfdisk -Y MBR --dump "$dev" || return
}

truncate -s0 whole.img
for f in part*.conf
do
        f=$(readlink -e _build/"${f%.conf}") || break
        ficlonerange.py "$f" whole.img
done

$sudo losetup -L -f whole.img
whole=$(losetup -j whole.img -O NAME --noheadings)
(pee_on_table "$whole") || exit

$sudo kpartx -u "$whole"

target=x86_64-efi
target=i386-pc
set -- grub-install --target=$target --removable

for f in part*.conf
do
        [ "$f" != part0.conf ] || continue
        notice "Processing $f"
        img=${f%.conf}
        dev=/dev/mapper/${whole#/dev/}p${img#part}
        mnt=${f%.conf}.mnt

        read_config_file || warn "Received error return from command: read_config_file $img"
        require_var name
        require_var type
        case "$type" in
                dm-verity-hashes|partition-table) ;;
                efi-system-partition|bios-grub)
                        mkdir -p "$mnt"
                        loudly $sudo mount "$dev" "$mnt"

                        if [ "$type" = 'bios-grub' ]
                        then
                                $sudo mkdir "$mnt"/grub
                                set -- "$@" --boot-directory="$mnt"
                        else
                                set -- "$@" --efi-directory="$mnt"
                        fi
                ;;
                *) mkdir -p "$mnt"
                   loudly $sudo mount "$dev" "$mnt"
                   ;;
                *) notice "Not mounting $name" ;;
        esac
done

loudly $sudo "$@" "$whole"

for f in part*.conf
do
        mnt=${f%.conf}.mnt
        if mountpoint -q "$mnt"
        then
                loudly $sudo umount "$mnt"
                $sudo rmdir "$mnt"
        fi
done
loudly $sudo kpartx -d "$whole"
loudly $sudo losetup -d "$whole"