#!/bin/sh
printHelp(){
    cat << EOF
Create a bootable media from a Windows iso image
Usage:
  $(basename "$0") -i <iso-image> -d <device> [OPTIONS]

  -i, --iso-image FILE      iso image to be used in the media creation
  -d, --device    DEVICE    device where the image must be burnt

  Optional parameters:

  -h, --help                display a help message like this
  -v, --verbose             turn on verbose output
  -V, --version             print molotov version and exit
  -u, --uefi                make an uefi-bootable media
                            (default is legacy)

Mandatory or optional arguments to long options are also mandatory 
or optional for any corresponding short options.

Report bugs to <https://codeberg.org/cizordj/molotov>

EOF
exit 0
}

printVersion(){
    cat << EOF
Molotov 1.3

Copyright (C) Cézar Augusto de Campos
License: GPL-3+
EOF
exit 0
}

handleArguments(){
    test ! $# -gt 0
    THERE_IS_NO_ARGUMENT=$?
    VERBOSE=1
    UEFI=1

    while true
    do
        if [ $THERE_IS_NO_ARGUMENT -eq 0 ]; then
            printHelp
        elif [ "$1" = "-h" ] || [ "$1" = "help" ] || [ "$1" = "--help" ]; then
            printHelp
        elif [ "$1" = "-i" ] || [ "$1" = "iso-image" ] || [ "$1" = "--iso-image" ]; then
            shift 1
            ISO_IMAGE="$1"
        elif [ "$1" = "-d" ] || [ "$1" = "device" ] || [ "$1" = "--device" ]; then
            shift 1
            TARGET_DEVICE="$1"
        elif [ "$1" = "-u" ] || [ "$1" = "uefi" ] || [ "$1" = "--uefi" ]; then
            UEFI=0
            printOut "Uefi mode enabled\n"
        elif [ "$1" = "-V" ] || [ "$1" = "version" ] || [ "$1" = "--version" ]; then
            printVersion
        elif [ "$1" = "-v" ] || [ "$1" = "verbose" ] || [ "$1" = "--verbose" ]; then
            VERBOSE=0
        elif [ $# -gt 0 ]; then
            printErr "Unknown option: $1\n"
            exit 64
        else
            break
        fi
        shift 1
    done

    if [ ! -e "$ISO_IMAGE" ] || [ -z "$ISO_IMAGE" ]; then
        printErr "Iso image file not found: $ISO_IMAGE\n"
        exit 66
    elif [ ! -e "$TARGET_DEVICE" ]; then
        printErr "The specified device was not found: $TARGET_DEVICE\n"
        exit 67
    elif [ ! -w "$TARGET_DEVICE" ]; then
        printErr "The specified device cannot be written by the current user: $TARGET_DEVICE\n"
        exit 68
    fi
}

printErr(){
    >&2 printf '%b' "$@"
}

umountEverything(){
    printOut "Umounting $TARGET_DEVICE and $ISO_IMAGE\n"
    umount -q "$ISO_IMAGE" "$TARGET_DEVICE"*
    printOut "Cleaning folders $ISO_MOUNTPOINT and $DEVICE_MOUNTPOINT\n"
    rmdir "$ISO_MOUNTPOINT" "$DEVICE_MOUNTPOINT" 2> /dev/null
}

printOut(){
    if [ $VERBOSE -eq 0 ]; then
        printf '%b' "$@"
    fi
}

formatDeviceForUefi(){
    printOut "Removing file system signature from $TARGET_DEVICE\n"
    wipefs -f -q -a "$TARGET_DEVICE"*
    printOut "Writing new partition table for $TARGET_DEVICE\n"
    sfdisk -q "$TARGET_DEVICE" << EOF
label: gpt
device: $TARGET_DEVICE
unit: sectors
sector-size: 512

$FIRST_PARTITION : start=        2048, size=       63488, type=EBD0A0A2-B9E5-4433-87C0-68B6B72699C7
$SECOND_PARTITION : start=       65536, type=EBD0A0A2-B9E5-4433-87C0-68B6B72699C7
EOF
    printOut "Formatting $FIRST_PARTITION with FAT32\n"
    mkfs.fat "$FIRST_PARTITION" || exit $?
    printOut "Formatting $SECOND_PARTITION with NTFS\n"
    mkfs.ntfs --label "$DISK_LABEL" --quiet --fast "$SECOND_PARTITION" || exit $?
}

formatDeviceForCsm(){
    printOut "Removing file system signature from $TARGET_DEVICE\n"
    wipefs -f -q -a "$TARGET_DEVICE"*
    printOut "Writing new partition table for $TARGET_DEVICE\n"
    sfdisk -q "$TARGET_DEVICE" << EOF
label: dos
device: $TARGET_DEVICE
unit: sectors
sector-size: 512
$FIRST_PARTITION : start= 2048, type=7, bootable
EOF
    printOut "Formatting $FIRST_PARTITION with NTFS\n"
    mkfs.ntfs --label "$DISK_LABEL" --fast --quiet "$FIRST_PARTITION" || exit $?
}

mountDevice(){
    printOut "Mounting $1 at $2\n"
    mkdir -p "$2"
    mount "$1" "$2"
}

dumpImageToDevice(){
    cp -r "$1"* "$2"
}

installBootloaderForCsm(){
    print_UUID(){
        blkid -o value -s UUID "$FIRST_PARTITION"
    }
    generate_config(){
        cat << EOF
set timeout=0
menuentry "Windows 10" {
    insmod ntfs
    insmod part_msdos
    insmod search
    insmod parttool
    insmod drivemap
    insmod ntldr
    search --no-floppy --fs-uuid --set=root $(print_UUID)
    parttool \${root} hidden-
    drivemap -s (hd0) \${root}
    ntldr /bootmgr
    boot
}
EOF
}
    BOOT_DIRECTORY="$DEVICE_MOUNTPOINT/boot"
    printOut "Installing grub on $TARGET_DEVICE\n"
    grub-install \
        --boot-directory="$BOOT_DIRECTORY" \
        --target="i386-pc" \
        "$TARGET_DEVICE"
    printOut "Generating grub configuration file\n"
    generate_config > "$BOOT_DIRECTORY/grub/grub.cfg"
}
installBootloaderForUefi(){
    print_UUID(){
        blkid -o value -s UUID "$SECOND_PARTITION"
    }
    generate_config(){
    cat << EOF
set timeout=0
menuentry "Windows 10" {
    insmod chain
    insmod search_fs_uuid
    search --set=root --fs-uuid $(print_UUID)
    chainloader (\$root)/efi/boot/bootx64.efi
}
EOF
    }
    BOOT_DIRECTORY="$DEVICE_MOUNTPOINT/boot"
    printOut "Installing grub in the flash drive's EFI partition\n"
    grub-install \
        --target="x86_64-efi" \
        --efi-directory="$DEVICE_MOUNTPOINT" \
        --boot-directory="$BOOT_DIRECTORY" \
        --removable \
        --no-nvram \
        --no-uefi-secure-boot \
        --install-modules="chain search_fs_uuid"
    printOut "Generating grub configuration file\n"
    generate_config > "$BOOT_DIRECTORY/grub/grub.cfg"
}

handleArguments "$@"

DEVICE_MOUNTPOINT="$(mktemp -d -u)"
DISK_LABEL="WINDOWS_1x"
ISO_MOUNTPOINT="$(mktemp -d -u)"
FIRST_PARTITION="${TARGET_DEVICE}1"
SECOND_PARTITION="${TARGET_DEVICE}2"

umountEverything

if [ $UEFI -eq 0 ]; then
    formatDeviceForUefi
    mountDevice "$SECOND_PARTITION" "$DEVICE_MOUNTPOINT"
    mountDevice "$ISO_IMAGE" "$ISO_MOUNTPOINT"
    dumpImageToDevice "$ISO_MOUNTPOINT/" "$DEVICE_MOUNTPOINT/"
    umountEverything
    mountDevice "$FIRST_PARTITION" "$DEVICE_MOUNTPOINT"
    installBootloaderForUefi
    umountEverything
    exit 0
else
    formatDeviceForCsm
    mountDevice "$FIRST_PARTITION" "$DEVICE_MOUNTPOINT"
    mountDevice "$ISO_IMAGE" "$ISO_MOUNTPOINT"
    dumpImageToDevice "$ISO_MOUNTPOINT/" "$DEVICE_MOUNTPOINT/"
    installBootloaderForCsm
    umountEverything
    exit 0
fi
