Check that ISO file and device exist

Fixes #1

README.md

@@ -21,6 +21,8 @@ Download latest [portable AppImage version](https://github.com/ValdikSS/windows2
 
 The program prints removable storage list if no arguments supplied.
 
+If you don't want to use AppImage, you'll need to install all dependencies and download [uefi-ntfs.img](https://github.com/pbatard/rufus/tree/master/res/uefi) from Rufus project.
+
 ### BIOS Boot
 
 BIOS Boot (Legacy Boot/UEFI-CSM) uses stock Windows 7 MBR and NTFS bootloader, courtesy of [ms-sys](http://ms-sys.sourceforge.net/) project.

windows2usb

@@ -24,6 +24,20 @@ function check_requirements() {
     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))}'
@@ -39,8 +53,12 @@ function format_drive() {
     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 > 429496296) {exit 1}; exit}'
+    7z l "$1" | awk '($6 ~ /install.wim/) {if ($4 > 4294967296) {exit 0}; exit 1}'
 }
 
 function create_partitions() {
@@ -68,6 +86,26 @@ function create_partitions() {
     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() {
     cat "$dirpath/uefi-ntfs.img" > "$1"
 }
@@ -79,9 +117,19 @@ function extract_iso() {
     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() {
-    umount "$1" || true
-    rm -r "$1"
+    if [ -d "$1" ];
+    then
+        umount "$1" || true
+        rm -r "$1"
+    fi
 }
 
 function sigint_handler() {
@@ -91,6 +139,8 @@ function sigint_handler() {
 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)."
@@ -108,11 +158,24 @@ if [[ "$EUID" == "0" ]];
 then
     dev="$1"
     isopath="$2"
+    isolabel=""
     labeltype="$3"
-    partpath="$(mktemp -d /run/winiso-write.XXXXXXXXXX)"
+    uefimode=
+    [[ "$labeltype" == "gpt" ]] && uefimode=1
+    [[ "$labeltype" == "gptntfs" ]] && uefimode=1
 
     check_requirements
-    trap sigint_handler INT
+    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" ]];
@@ -122,14 +185,14 @@ then
 
         echo "${bold} == Creating NTFS partition ==${normal}"
         create_partitions "dos" "$dev"
-        mkfs.ntfs -f "${dev}1"
+        mkfs.ntfs -L "$isolabel" -f "$(get_dev_partition_num "${dev}" "1")"
 
         echo "${bold} == Writing bootloader ==${normal}"
         ms-sys -7 "${dev}"
-        ms-sys -n "${dev}1"
+        ms-sys -n "$(get_dev_partition_num "${dev}" "1")"
 
         echo "${bold} == Mounting data partition ==${normal}"
-        mount "${dev}1" "$partpath"
+        mount "$(get_dev_partition_num "${dev}" "1")" "$partpath"
 
     # GPT
     elif [[ "$labeltype" == "gpt" ]];
@@ -138,6 +201,7 @@ then
         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
@@ -146,10 +210,10 @@ then
 
         echo "${bold} == Creating UEFI FAT partition ==${normal}"
         create_partitions "gpt" "$dev"
-        mkfs.vfat "${dev}1"
+        mkfs.vfat -n "${isolabel:0:11}" "$(get_dev_partition_num "${dev}" "1")"
 
         echo "${bold} == Mounting data partition ==${normal}"
-        mount "${dev}1" "$partpath"
+        mount "$(get_dev_partition_num "${dev}" "1")" "$partpath"
 
     # GPT FAT32 + NTFS
     elif [[ "$labeltype" == "gptntfs" ]];
@@ -159,16 +223,29 @@ then
 
         echo "${bold} == Creating UEFI FAT and Microsoft NTFS partitions ==${normal}"
         create_partitions "gptntfs" "$dev"
-        write_uefintfs "${dev}1"
-        mkfs.ntfs -f "${dev}2"
+        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 "${dev}2" "$partpath"
+        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"
 
@@ -177,12 +254,16 @@ else
     if [[ "$APPIMAGE" ]];
     then
         scriptpath="$APPIMAGE"
-    fi;
+    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" "$@"
+            "$privesc" "$scriptpath" "$PARAM1" "$PARAM2" "$@"
             exit
         fi
     done