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

: ${GRUB_CONFIG:=../conf/grub.cfg}
samizdat_linux_dir=/

if [ ! "$GPG_INPUT_DIR" ]
then
    for d in $GPG_INPUT_DIR /root/.gnupg /cdrom/gnupghome
    do
        $sudo [ -d "$d" ] || continue
        GPG_INPUT_DIR=$d
        break
    done
fi


add_initrd()
{
    initrd_suffix=.samizdat
    $sudo mkdir -p "$mnt"/linux
    $sudo rsync  -aL --info=STATS "${1}vmlinuz${2}" "$mnt"/linux/vmlinuz
    $sudo rsync  -aL --info=STATS "${1}initrd.img${2}${2:+$initrd_suffix}" "$mnt"/linux/initrd.img
}

add_grub_cfg()
{
    $sudo mkdir -p "$mnt"/grub
    $sudo cp -aL "$GRUB_CONFIG" "$mnt"/grub
}

install_boot_dir()
{
    local mnt="$1"
    add_grub_cfg
    add_initrd "$samizdat_linux_dir"/ "${version_suffix}"
}

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
                        bios-grub)            ;;
                        efi-system-partition) mkfs.fat -F 32 -I "$tmp" || die "mkfs.vfat failed" ;;
                        boot|samizdat-keys)   mkfs.btrfs -q "$tmp" || die "mkfs.btrfs failed" ;;
                        dm-verity-data)
                                partuuid=${root_hash:0:32}
                                set_var partuuid
                                cp -f -T --reflink "$data_path" "$builddir"/"$partuuid"
                                ln -sfT "$partuuid" "$tmp"
                        ;;
                        dm-verity-hashes)
                                partuuid=${root_hash:32:32}
                                set_var partuuid
                                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|boot|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}"
}

format_guid()
{
    set -- "$(tr -dc 0-9a-fA-F <<< "$1")"
    printf '%s\n' "${1:0:8}-${1:8:4}-${1:12:4}-${1:16:4}-${1:20:12}"
}

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

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

        start_sectors=$((start / 512))
        size_sectors=$((devsz / 512))

        typecode=0FC63DAF-8483-4772-8E79-3D69D8477DE4
        case "$type" in
                partition-table) start=$((start + devsz)); return;;
                efi-system-partition) typecode=C12A7328-F81F-11D2-BA4B-00A0C93EC93B ;;
                dm-verity-data|dm-verity-hashes)
                    check_for_key_only_skip || return 0
                    case "$name" in
                        samizdat-rootfs|samizdat-root-patch) typecode=4f68bce3-e8cd-4db1-96e7-fbcaf984b709 ;;
                        samizdat-root-patch-verity) typecode=2c7357ed-ebd2-46d9-aec1-23d437ec2bf5 ;;
                        samizdat-root-seed) ;; # keep default
                        samizdat-root-seed-verity) ;; # keep default
                    esac
                    ;;
                root) typecode=4f68bce3-e8cd-4db1-96e7-fbcaf984b709 ;;
                bios-grub)
                    typecode=21686148-6449-6E6F-744E-656564454649
                    printf 'start=1, size=%d, type=ee\n' \
                           "$((start_sectors - 1))" \
                           >> "$DOS_TABLE_FILE"
                    printf 'start=%d, size=%d, type=83, bootable\n' \
                           "$start_sectors" \
                           "$size_sectors" \
                           >> "$DOS_TABLE_FILE"
                    printf 'start=%d, type=ee\n' \
                           "$((start_sectors + size_sectors))" \
                           >> "$DOS_TABLE_FILE"
                    ;;
        esac

        (
            exec >> "$GPT_TABLE_FILE"
            printf '%d: start=%d, size=%d, type=%s, name="%s"' \
                    "$i" \
                    "$start_sectors" \
                    "$size_sectors" \
                    "$typecode" \
                    "$name"
            if [ "$partuuid" ]
            then
                printf ', uuid=%s' "$(format_guid "$partuuid")"
            fi
            echo
        )
        let ++i
        start=$((start + devsz))
}

KEY_ONLY=
if [ "$1" = 'key' ]
then
        KEY_ONLY=y
fi

set -e

builddir=_build
mkdir -p "$builddir"

iterate_partitions build_partition_image

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

if [ "$KEY_ONLY" ]
then
        target=key.img
else
        target=whole.img
fi

if whole=$(losetup -j "$target" -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()
{
        GPT_TABLE_FILE=$builddir/table.gpt
        DOS_TABLE_FILE=$builddir/table.mbr

        local dev="$1"
        i=1
        start=0

        sfdisk_init "$dev" > "$GPT_TABLE_FILE"
        : > "$DOS_TABLE_FILE"
        iterate_partitions create_ptable_conf

        set -x
        $sudo sfdisk "$dev" < "$GPT_TABLE_FILE" || return
        $sudo sfdisk -Y dos "$dev" < "$DOS_TABLE_FILE" || return
}

check_for_key_only_skip()
{
        [ "$KEY_ONLY" ] || return 0
        case "$type" in
                dm-verity-data|dm-verity-hashes) false ;;
                *) true ;;
        esac
}

clone_parts_to_target()
{

        check_for_key_only_skip || return 0
        f=$(readlink -e _build/"${f%.conf}") || return
        ficlonerange.py "$f" "$target"
}

truncate -s0 "$target"
iterate_partitions clone_parts_to_target

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

$sudo kpartx -u "$whole"

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|bios-grub) ;;
                efi-system-partition|boot|samizdat-keys)
                        mkdir -p "$mnt"
                        loudly $sudo mount "$dev" "$mnt"

                        case "$type" in
                            boot)
                                BOOT_DIR=$mnt
                                install_boot_dir "$mnt"
                                ;;
                            samizdat-keys)
                                $sudo rsync  -a  --info=STATS  "$GPG_INPUT_DIR"/ "$mnt"/gnupghome/
                                ;;
                            efi-system-partition)
                                EFI_DIR=$mnt
                                ;;
                        esac
                ;;
                *) notice "Not mounting $name" ;;
        esac
done

loudly $sudo eatmydata -- grub-install --target=i386-pc    --recheck                --boot-directory="$BOOT_DIR" "$whole"
loudly $sudo eatmydata -- grub-install --target=x86_64-efi --recheck   --removable   --efi-directory="$EFI_DIR"  "$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"