#!/usr/bin/env bash
set -e # errexit
set -f # noglob
set +o braceexpand # disable braceexpand
shopt -s nocasematch # for "if labeltype"

export LANG=C

bold="$(tput bold)"
normal="$(tput sgr0)"
scriptpath="$(readlink -f -- "$0")"
dirpath="$(dirname "$scriptpath")"


function check_requirements() {
    local reqs=(awk lsblk 7z mkfs.vfat mkfs.ntfs sfdisk ms-sys mktemp)

    for req in ${reqs[*]}; do
        if ! command -v "$req" > /dev/null;
        then
            echo "${bold} == ERROR: $req is required but not found ==${normal}"
            exit 104
        fi
    done
}

function check_iso_and_device() {
    if [ ! -f "$1" ]
    then
        echo "${bold} == ERROR: ISO file not found! ==${normal}"
        exit 106
    fi

    if [ ! -b "$2" ]
    then
        echo "${bold} == ERROR: Device file not found! ==${normal}"
        exit 107
    fi
}

function list_removable_drives() {
    lsblk -d -I 8 -o RM,NAME,SIZE,MODEL | \
        awk '($1 == 1) {print "/dev/" substr($0, index($0, $2))}'
}

function format_drive() {
    if [[ "$1" == "dos" ]] || [[ "$1" == "gpt" ]];
    then
        echo "label: $1" | sfdisk "$2"
    else
        echo "${bold} == format_drive: INTERNAL ERROR ==${normal}"
        exit 101
    fi
}

function get_iso_name() {
    7z l "$1" | awk '/Comment = / {print $3; exit 0}'
}

function check_installwim_gt4gib() {
    7z l "$1" | awk '($6 ~ /install.wim/) {if ($4 > 4294967296) {exit 0}; exit 1}'
}

function create_partitions() {
    local mbrscript="- - 7 *"
    local gptscript="- - EBD0A0A2-B9E5-4433-87C0-68B6B72699C7 *"
    local gptntfsscript="- 1MiB U *\n- - EBD0A0A2-B9E5-4433-87C0-68B6B72699C7 *"

    if [[ "$1" == "dos" ]] || [[ "$1" == "gpt" ]] || [[ "$1" == "gptntfs" ]];
    then
        if [[ "$1" == "dos" ]];
        then
            echo -e "$mbrscript" | sfdisk "$2"

        elif [[ "$1" == "gpt" ]];
        then
            echo -e "$gptscript" | sfdisk "$2"

        elif [[ "$1" == "gptntfs" ]];
        then
            echo -e "$gptntfsscript" | sfdisk "$2"
        fi
    else
        echo "${bold} == create_partitions: INTERNAL ERROR ==${normal}"
        exit 102
    fi
}

function get_dev_partition_num() {
    # $1 - device
    # $2 - partition number

    if [ -b "${1}${2}" ];
    then
        echo "${1}${2}"
        return
    fi

    if [ -b "${1}p${2}" ];
    then
        echo "${1}p${2}"
        return
    fi

    echo "${bold} == get_dev_partition_num: INTERNAL ERROR ==${normal}"
    exit 105
}

function write_uefintfs() {
    if [ -a "$dirpath/uefi-ntfs.img" ]; then
        cat "$dirpath/uefi-ntfs.img" > "$1"
    elif [ -a "/usr/share/windows2usb/uefi-ntfs.img" ]; then
        cat "/usr/share/windows2usb/uefi-ntfs.img" > "$1"
    fi
}

function extract_iso() {
    # $1 - isopath
    # $2 - destdir

    7z x "$1" -o"$2"
}

function extract_bootmgfw_from_installwim() {
    # $1 - install.wim path
    # $2 - destdir

    7z e "$1" -o"$2" 1/Windows/Boot/EFI/bootmgfw.efi
}

function umount_rm_path() {
    if [ -d "$1" ];
    then
        umount "$1" || true
        rm -r "$1"
    fi
}

function sigint_handler() {
    umount_rm_path "$partpath"
}

if [[ ! "$3" ]];
then
    echo "Windows 7/8/8.1/10 ISO to Flash Drive burning utility"
    echo "WARNING: this program will delete all existing data on your drive!"
    echo
    echo "$(basename "$0")" "<device> <windows iso> <mbr/gpt/gptntfs>"
    echo
    echo "Use MBR mode for old computers with BIOS (without UEFI), or for UEFI legacy mode (CSM)."
    echo "Use GPT mode for new UEFI computers."
    echo "Use 'gptntfs' mode for custom ISO images with install.wim file greater than 4 GiB."
    echo "NOTE: gptntfs does not support Secure Boot!"
    echo
    echo "${bold} == Removable storage list ==${normal}"

    list_removable_drives
    exit
fi

if [[ "$EUID" == "0" ]];
then
    dev="$1"
    isopath="$2"
    isolabel=""
    labeltype="$3"
    uefimode=
    [[ "$labeltype" == "gpt" ]] && uefimode=1
    [[ "$labeltype" == "gptntfs" ]] && uefimode=1

    check_requirements
    check_iso_and_device "$isopath" "$dev"

    isolabel="$(get_iso_name "$isopath")"
    if [ $? -ne 0 ]
    then
        isolabel=""
    fi
    echo "${bold} == Working with ISO $isolabel ==${normal}"

    partpath="$(mktemp -d /run/windows2usb.XXXXXXXXXX)"
    trap sigint_handler INT EXIT

    # MBR
    if [[ "$labeltype" == "mbr" ]];
    then
        echo "${bold} == Creating new MBR-formatted partition table ==${normal}"
        format_drive "dos" "$dev"

        echo "${bold} == Creating NTFS partition ==${normal}"
        create_partitions "dos" "$dev"
        mkfs.ntfs -L "$isolabel" -f "$(get_dev_partition_num "${dev}" "1")"

        echo "${bold} == Writing bootloader ==${normal}"
        ms-sys -7 "${dev}"
        ms-sys -n "$(get_dev_partition_num "${dev}" "1")"

        echo "${bold} == Mounting data partition ==${normal}"
        mount "$(get_dev_partition_num "${dev}" "1")" "$partpath"

    # GPT
    elif [[ "$labeltype" == "gpt" ]];
    then
        if check_installwim_gt4gib "$isopath";
        then
            echo "${bold} == ERROR: install.wim is greater than 4 GiB ==${normal}"
            echo "${bold} == ERROR: Please use gptntfs mode ==${normal}"
            echo "NOTE: gptntfs does not support Secure Boot!"
            umount_rm_path "$partpath"
            exit 103
        fi
        echo "${bold} == Creating new GPT-formatted partition table ==${normal}"
        format_drive "gpt" "$dev"

        echo "${bold} == Creating UEFI FAT partition ==${normal}"
        create_partitions "gpt" "$dev"
        mkfs.vfat -n "${isolabel:0:11}" "$(get_dev_partition_num "${dev}" "1")"

        echo "${bold} == Mounting data partition ==${normal}"
        mount "$(get_dev_partition_num "${dev}" "1")" "$partpath"

    # GPT FAT32 + NTFS
    elif [[ "$labeltype" == "gptntfs" ]];
    then
        echo "${bold} == Creating new GPT-formatted partition table ==${normal}"
        format_drive "gpt" "$dev"

        echo "${bold} == Creating UEFI FAT and Microsoft NTFS partitions ==${normal}"
        create_partitions "gptntfs" "$dev"
        write_uefintfs "$(get_dev_partition_num "${dev}" "1")"
        mkfs.ntfs -L "$isolabel" -f "$(get_dev_partition_num "${dev}" "2")"

        echo "${bold} == Mounting data partition ==${normal}"
        mount "$(get_dev_partition_num "${dev}" "2")" "$partpath"
    fi

    echo "${bold} == Extracting files from ISO to the partition ==${normal}"
    extract_iso "$isopath" "$partpath"

    if [[ "$uefimode" ]];
    then
        if [ ! -f "$partpath/efi/boot/bootx64.efi" ] && \
           [ ! -f "$partpath/efi/boot/bootia32.efi" ];
        then
            echo "${bold} == Extracting UEFI bootloader from install.wim ==${normal}"
            mkdir -p "$partpath/efi/boot/" || true
            extract_bootmgfw_from_installwim "$partpath/sources/install.wim" "$partpath/efi/boot/"
            mv "$partpath/efi/boot/bootmgfw.efi" "$partpath/efi/boot/bootx64.efi"
            cp "$partpath/efi/boot/bootx64.efi" "$partpath/efi/boot/bootia32.efi"
        fi
    fi

    echo "${bold} == Unmounting partition ==${normal}"
    umount_rm_path "$partpath"

    echo "${bold} == All done! ==${normal}"
else
    if [[ "$APPIMAGE" ]];
    then
        scriptpath="$APPIMAGE"
    fi

    [ -f "$1" ] && PARAM1="$(readlink -f -- "$1")" || PARAM1="$1"
    [ -f "$2" ] && PARAM2="$(readlink -f -- "$2")" || PARAM2="$2"
    shift; shift;

    privescs=(pkexec sudo)
    for privesc in ${privescs[*]}; do
        if command -v "$privesc" > /dev/null; then
            "$privesc" "$scriptpath" "$PARAM1" "$PARAM2" "$@"
            exit
        fi
    done

    echo "${bold} == ERROR: no pkexec or sudo found. ==${normal}"
    echo "${bold} == ERROR: please run this script as root manually. ==${normal}"
fi