winusb

#!/usr/bin/env bash
# shellcheck disable=SC2026
# SC2026 disabled because of false positives(ShellCheck issue #923)

# A Linux program to create bootable Windows USB stick from a real Windows DVD or an image
# Copyright © 2013 Colin GILLE / congelli501
# Copyright © 2017 slacka et. al.

# This file is part of WinUSB.
#
# WinUSB is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# WinUSB is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with WinUSB  If not, see <http://www.gnu.org/licenses/>.

# Notes:
#
# * Use --strip instead of --no-symlinks for compatibility of Ubuntu 14.04(EoL: April 2019)

## Makes debuggers' life easier - Unofficial Bash Strict Mode
## BASHDOC: Shell Builtin Commands - Modifying Shell Behavior - The Set Builtin
set -o errexit
set -o errtrace
set -o nounset
set -o pipefail

## Non-overridable Primitive Variables
## BASHDOC: Shell Variables » Bash Variables
## BASHDOC: Basic Shell Features » Shell Parameters » Special Parameters
if [ -v "BASH_SOURCE[0]" ]; then
	RUNTIME_EXECUTABLE_PATH="$(realpath --strip "${BASH_SOURCE[0]}")"
	RUNTIME_EXECUTABLE_FILENAME="$(basename "${RUNTIME_EXECUTABLE_PATH}")"
	RUNTIME_EXECUTABLE_NAME="${RUNTIME_EXECUTABLE_FILENAME%.*}"
	RUNTIME_EXECUTABLE_DIRECTORY="$(dirname "${RUNTIME_EXECUTABLE_PATH}")"
	RUNTIME_COMMANDLINE_BASECOMMAND="${0}"
	# Keep these partially unused variables declared
	# shellcheck disable=SC2034
	declare -r\
		RUNTIME_EXECUTABLE_PATH\
		RUNTIME_EXECUTABLE_FILENAME\
		RUNTIME_EXECUTABLE_NAME\
		RUNTIME_EXECUTABLE_DIRECTORY\
		RUNTIME_COMMANDLINE_BASECOMMAND
fi
declare -ar RUNTIME_COMMANDLINE_PARAMETERS=("${@}")

declare install_method
declare\
	source_media\
	target_media
declare -i\
	verbose='0'\
	only_for_gui='0'\
	no_color='0'
declare\
	command_mkdosfs\
	command_grubinstall
declare\
	source_fs_mountpoint
	target_fs_mountpoint
declare -i BLOCK_SIZE="((4 * 1024 * 1024))" # 4MiB
declare -i pulseCurrentPid=0 # FIXME: No documentation for this non-trivial parameter

init(){
	if ! process_commandline_parameters; then
		print_help
		exit 1
	fi

	if ! check_runtime_dependencies; then
		exit 1
	fi

	if [ ! "$(id --user)" = 0 ]; then
		printf --\
			"%s\n%s\n"\
			"Warning: You are not running WinUSB as root!"\
			"Warning: This might be the reason of the following failure." >&2
	fi

	if ! check_runtime_parameters; then
		exit 1
	fi

	declare\
		target_device\
		target_partition

	pulse on

	if [ "${install_method}" = 'edit' ]; then
		target_partition="${target_media}"
		target_device="$(echo "${target_media}" | sed "s/[0-9]*$//")"
		if [ "${verbose}" = '1' ]; then
			echo "Target device is '${target_device}'."
		fi

		# If target partition is mounted in the first place, bail out
		if [ "$(mount | grep -c "${target_partition}")" != 0 ]; then
			echo_with_color red "Error: target partition is currently mounted, unmount the partition then try again"
			exit 1
		fi

	else # install_method = format
		target_device="${target_media}"
		target_partition="${target_device}1"
		# If target media busy, bail out
		if [ "$(mount | grep -c "${target_device}")" != 0 ]; then
			echo_with_color red "Error: target device is currently busy, unmount all mounted partitions in target device then try again"
			exit 1
		fi
		destroy_and_rebuild_entire_disk
	fi

	mount_source_and_target_filesystem

	setup_traps

	copy_filesystem_files "${source_fs_mountpoint}" "${target_fs_mountpoint}"

	# As Windows 7's installation media doesn't place the required EFI bootloaders in the right location, we extract them from the system image manually
	# TODO: Functionize Windows 7 checking
	if ( grep -Eq "^MinServer=7[0-9]{3}\.[0-9]" "${source_fs_mountpoint}/sources/cversion.ini") && [ -f "${source_fs_mountpoint}/bootmgr.efi" ]; then
		echo_with_color yellow "Source media is Windows 7-based, applying workaround to make it support UEFI booting"

		if command -v "7z" >/dev/null 2>&1; then
			workaround_support_windows_7_uefi_boot
		else
			echo_with_color yellow "Warning: '7z' utility not found, workaround is not applied." >&2
		fi
	fi

	install_bootloader_grub

	cleanup_mountpoints "${source_fs_mountpoint}" "${target_fs_mountpoint}" 'normalCleanup'

	echo_with_color green "Done :-)  You may detach the usb storage now."

	pulse off

	exit 0
}; declare -fr init

install_bootloader_grub(){
	util_check_function_parameters_quantity "${FUNCNAME[0]}" 0 "${#}"

	echo_with_color blue "Installing GRUB bootloader for legacy boot support..."
	"${command_grubinstall}" --target=i386-pc --boot-directory="${target_fs_mountpoint}" --force "${target_device}"

	uuid=$(blkid -o value -s UUID "${partition}")

	echo_with_color blue "Installing grub.cfg..."
	grub_cfg="${target_fs_mountpoint}/grub/grub.cfg"
	mkdir -p "$(dirname "${grub_cfg}")"
	{
		cat <<- END_OF_FILE
			echo '------------------------------------'
			echo '|      Windows USB - Loading...    |'
			echo '------------------------------------'
			insmod fat
			search --no-floppy --fs-uuid $uuid --set root
			ntldr /bootmgr
			boot
		END_OF_FILE
	} > "${grub_cfg}"
}; declare -fr install_bootloader_grub

workaround_support_windows_7_uefi_boot(){
	util_check_function_parameters_quantity "${FUNCNAME[0]}" 0 "${#}"

	mkdir -p "${target_fs_mountpoint}/EFI/Boot"
	7z e -so "${source_fs_mountpoint}/sources/install.wim" "1/Windows/Boot/EFI/bootmgfw.efi" > "${target_fs_mountpoint}/EFI/Boot/bootx64.efi"
}; declare -fr workaround_support_windows_7_uefi_boot

mount_source_and_target_filesystem(){
	util_check_function_parameters_quantity "${FUNCNAME[0]}" 0 "${#}"

	echo_with_color blue "Mounting source and target filesystems..."
	source_fs_mountpoint="/media/winusb_source_$(date +%s)_$$"
	target_fs_mountpoint="/media/winusb_target_$(date +%s)_$$"

	mkdir -p "${source_fs_mountpoint}"
	if [ -f "${source_media}" ]; then # ${source_media} is an ISO image
		mount -o loop,ro -t udf,iso9660 "${source_media}" "${source_fs_mountpoint}"
	else # ${source_media} is a real DVD drive (block device)
		mount -o ro "${source_media}" "${source_fs_mountpoint}"
	fi

	mkdir -p "${target_fs_mountpoint}"
	mount "${target_partition}" "${target_fs_mountpoint}"

	free_space=$(df --block-size 1 "${target_fs_mountpoint}" | grep "${target_partition}" | awk '{print $4}')
	needed_space=$(du -s "${source_fs_mountpoint}" --bytes | awk '{print $1}')

	additional_space_required_for_grub_installation=((1000 * 1000 * 10)) # 10MiB
	((needed_space = needed_space + additional_space_required_for_grub_installation))

	if [ "${needed_space}" -gt "${free_space}" ]; then
		echo "Error: Not enough free space on '${target_partition}'!" >&2
		echo "Error: We required ${needed_space} bytes but '${target_partition}' only has ${free_space} bytes." >&2
		exit 1
	fi
}; declare -fr mount_source_and_target_filesystem

setup_traps(){
	util_check_function_parameters_quantity "${FUNCNAME[0]}" 0 "${#}"

	trap 'cleanup_mountpoints "${source_fs_mountpoint}" "${target_fs_mountpoint}" "error"' ERR
	trap 'cleanup_mountpoints "${source_fs_mountpoint}" "${target_fs_mountpoint}" "interrupt"' INT
	trap 'cleanup_mountpoints "${source_fs_mountpoint}" "${target_fs_mountpoint}" "kill"' TERM
	trap 'cleanup_mountpoints "${source_fs_mountpoint}" "${target_fs_mountpoint}" "exit"' EXIT
	trap 'cleanup_mountpoints "${source_fs_mountpoint}" "${target_fs_mountpoint}" "other signal"' HUP QUIT PIPE
}; declare -fr setup_traps

# Utility function to check if function parameters quantity is legal
# NOTE: non-static function parameter quantity(e.g. either 2 or 3) is not supported
util_check_function_parameters_quantity(){
	if [ "${#}" -ne 2 ]; then
		printf --\
			"%s: FATAL: Function requires %u parameters, but %u is given\n"\
			"${FUNCNAME[0]}"\
			2\
			"${#}"
		exit 1
	fi

	# The expected given quantity
	local -i expected_parameter_quantity="${1}"; shift
	# The actual given parameter quantity, simply pass "${#}" will do
	local -i given_parameter_quantity="${1}"

	if [ "${given_parameter_quantity}" -ne "${expected_parameter_quantity}" ]; then
		printf --\
			"%s: FATAL: Function requires %u parameters, but %u is given\n"\
			"${FUNCNAME[1]}"\
			"${expected_parameter_quantity}"\
			"${given_parameter_quantity}"\
			1>&2
		exit 1
	fi
	return 0
}; declare -fr util_check_function_parameters_quantity

destroy_and_rebuild_entire_disk(){
	local storage_device="${target_media}"

	# Create new PC, a.k.a. MBR, a.k.a. msdos style partition table(and overwrite the old one, whatever it was)
	parted --script\
		"${storage_device}"\
		mklabel\
		msdos

	# Create partition
	# We start at 4MiB for grub (it needs a post-mbr gap for its code) and alignment of flash memery block erase segment in general, for details see http://www.gnu.org/software/grub/manual/grub.html#BIOS-installation and http://lwn.net/Articles/428584/
	parted --script\
		"${storage_device}"\
		mkpart\
		primary\
		fat32\
		4MiB\
		-- -1s # last sector of the disk

	# Reload the new partition table
	# FIXME: Is `blockdev` and `partprobe` both required?
	blockdev\
		--rereadpt\
		"${storage_device}"\
		|| true
	partprobe\
		"${storage_device}"
	echo "Wait 3 seconds for block device nodes to populate..."
	sleep 3

	# get first partition
	local partition
	partition="${storage_device}1"

	# Create filesystem on partition
	"$command_mkdosfs"\
		-F 32\
		-n 'Windows USB' \
		"${partition}"
}; declare -fr destroy_and_rebuild_entire_disk

check_runtime_dependencies(){
	local result="success"

	if which 'mkdosfs' > /dev/null; then
		command_mkdosfs='mkdosfs'
	elif which 'make.msdos' > /dev/null; then
		command_mkdosfs='mkfs.msdos'
	else
		echo_with_color red\
			"${FUNCNAME[0]}: Error: mkdosfs or mkfs.msdos command not found!" >&2
		echo_with_color red\
			"${FUNCNAME[0]}: Error: Please make sure that dosfstools is properly installed!" >&2
		result="failed"
	fi

	if which 'grub-install' > /dev/null; then
		command_grubinstall='grub-install'
	elif which 'grub2-install' > /dev/null; then
		command_grubinstall='grub2-install'
	else
		echo_with_color red "${FUNCNAME[0]}: Error: grub-install or grub2-install command not found!" >&2
		echo_with_color red "${FUNCNAME[0]}: Error: Please make sure that GNU GRUB is properly installed!" >&2
		result="failed"
	fi

	if [ "${result}" == "failed" ]; then
		return 1
	else
		return 0
	fi
}; declare -fr check_runtime_dependencies

check_runtime_parameters(){
	local -i result="0" # success

	if [ -z "${install_method}" ]; then
		echo "Error: No install method specified!" >&2
		echo "" >&2
		print_help
		result="1"
	fi

	if [ ! -f "${source_media}" ] \
		&& [ ! -b "${source_media}" ]; then
		echo "Error: source media '${source_media}' not found or not a regular file or a block device!" >&2
		result="1"
	fi

	if ! [ -b "${target_media}" ]; then
		echo "Error: the device \"${target_media}\" is not special block!" >&2
		result="1"
	fi

	return "${result}"
}; declare -fr check_runtime_parameters

pulse(){
	if [ "$only_for_gui" -eq 1 ]; then
		if [ ! "$pulseCurrentPid" -eq 0 ]; then
			kill "$pulseCurrentPid"
			#wait "$pulseCurrentPid"
			pulseCurrentPid=0
		fi

		if [ "$1" = 'on' ]; then
			cd /; while true; do sleep 0.05; echo 'pulse'; done &
			pulseCurrentPid="$!"
		elif [ "$1" != 'off' ]; then
			echo "Pulse: bad argument" >&2
			exit 1
		fi
	fi
}; declare -fr pulse

# Configure the terminal to print future messages with certain color
# $1:
#     <color>: color of the next message
#     none: reset color to default
switch_terminal_text_color()
{
	case "$1" in
		none)
			tput sgr0
			;;
		black)
			echo -en "\033[0;30m"
			;;
		red)
			echo -en "\033[0;31m"
			;;
		green)
			echo -en "\033[0;32m"
			;;
		yellow)
			echo -en "\033[0;33m"
			;;
		blue)
			echo -en "\033[0;34m"
			;;
		white)
			echo -en "\033[0;37m"
			;;
	esac
}; declare -fr switch_terminal_text_color

# Print message with color
echo_with_color(){
	util_check_function_parameters_quantity 2 "${#}"
	local -r message_color="${1}"; shift
	local -r message_body="${1}"

	if [ "${no_color}" -eq 1 ]; then
		echo -e "${message_body}"
	else
		switch_terminal_text_color "${message_color}"
		echo -e "${message_body}"
		switch_terminal_text_color none
	fi
}

# Unmount mounted filesystems and clean-up mountpoints before exiting program
# $1: Windows ISO mountpoint
# $2: Target mountpoint
# $3: Exit reason
cleanup_mountpoints(){
	util_check_function_parameters_quantity 3 "${#}"

	set +o errexit # We need to clean up everything we can

	trap - HUP QUIT PIPE INT TERM EXIT

	pulse off

	printf "\n" # Proactively separate from previous command output
	if [ "$3" == 'kill' ]; then
		echo_with_color yellow "Terminated by external signal!"  >&2
	elif [ "$3" == 'error' ]; then
		echo_with_color red "Internal error occurred!"  >&2
	elif [ "$3" == 'interrupt' ]; then
		echo_with_color yellow "Interrupted by user"
	else
		echo_with_color green "Exiting..."
	fi

	echo "Synchronizing cached writes to persistent storage..."

	pulse on
	sync
	pulse off

	echo "Cleaning..."

	cd /
	sync
	wait 2> /dev/null
	sleep 1

	echo_with_color green "Unmounting and removing '$1'..."
	# shellcheck disable=2015
	umount "$1" && rm -rf "$1" || true

	echo_with_color green "Unmounting and removing '$2'..."
	# shellcheck disable=2015
	umount "$2" && rm -rf "$2" || true

	if [ "$3" = 'normalCleanup' ]; then
		exit 0
	else
		exit 2
	fi
}

hideNoErr(){
	set +o errexit
	TEXT="$( "$@" 2>&1 )"
	RETURN="$?"
	set -o errexit

	if [ "$RETURN" -ne 0 ]; then
		echo_with_color red "Error while running command '$*' (exit code : $RETURN)!" >&2
		echo_with_color red "Read command output (Y/n)?" >&2

		rep='y'
		read -r rep || true # If we can't read, don't quit

		if [ "${rep}" != "n" ] && [ "${rep}" != "N" ]; then
			echo "$TEXT"
		fi

		echo_with_color red "Press ENTER to exit (tmpDir will be removed)" >&2
		read -r
	fi

	return "$RETURN"
}

# Copy source_file to dest_file, overwrite file if dest_file exists
copy_file(){
	util_check_function_parameters_quantity 3 "${#}"
	source_file="$1"
	dest_file="$2"
	whole_current_size="$3"

	# Calculate block count of the source file
	size=$(stat --format=%s "${source_file}")
	((stepNb = size / BLOCK_SIZE + 1))

	if [ -f "${dest_file}" ]; then
		rm "${dest_file}"
	fi

	# Copy file block by block
	i=0
	while [ "$i" -lt "$stepNb" ]; do
		dd if="${source_file}" bs="${BLOCK_SIZE}" skip="$i" seek="$i" of="${dest_file}" count=1 2> /dev/null
		((i = i + 1))

		# Send current copied size
		local copied_size
		copied_size=$(stat --format=%s "${dest_file}")
		((copied_size_total = copied_size + whole_current_size)) || true
		#echo "$copied_size_total"

		((percent = (copied_size_total * 100) / totalSize)) || true
		echo -en "${percent}%\r"
	done

	return 0
}; declare -fr copy_file

# Copying all files from one filesystem to another, with progress reporting
# $1: Source filesystem's mountpoint
# $2: Target filesystem's mountpoint
copy_filesystem_files(){
	util_check_function_parameters_quantity 2 "${#}"

	# FIXME: Why do we `pulse off` and on here?
	pulse off

	echo_with_color blue "Copying files from Windows installation media..."
	totalSize=$(du -s "$1" --bytes | awk '{print $1}')
	currentSize=0

	targetDir="$2"
	mkdir -p "$targetDir"
	targetDir=$(readlink -f "$targetDir")

	cd "$1"
	file_list=$(find .)

	while read -r source_file; do
		if [ "${source_file}" != '.' ]; then
			dest_file="$targetDir/${source_file}"

			file_size=$(stat --format=%s "${source_file}")

			if [ -d "${source_file}" ]; then
				mkdir -p "${dest_file}"
			elif [ -f "${source_file}" ]; then
				if [ "${file_size}" -lt "${BLOCK_SIZE}" ]; then
					cp "${source_file}" "${dest_file}"
				else
					copy_file "${source_file}" "${dest_file}" "$currentSize"
				fi
			else
				echo "Error: Unknown type of '${source_file}'!" >&2
				exit 1
			fi

			# Progress
			((currentSize = currentSize + file_size)) || true
			((percent = (currentSize * 100) / totalSize)) || true
			echo -en "${percent}%\r"
		fi
	done < <(echo "$file_list")

	pulse on

	return 0
}

print_help(){
	util_check_function_parameters_quantity 0 "${#}"

	echo -e "${RUNTIME_EXECUTABLE_NAME} usage"
	echo -e "Install a Windows ISO on an FAT partition and edit MBR of the device"
	echo -e "\t""${RUNTIME_EXECUTABLE_NAME} --install <iso path> <partition>"
	echo -e "\t""Example: ${RUNTIME_EXECUTABLE_NAME} --install win7_amd64.iso /dev/sdd1"
	echo -e
	echo -e "Completely format a drive and install the ISO on it"
	echo -e "\t""${RUNTIME_EXECUTABLE_NAME} --format <iso path> <device>"
	echo -e "\t""Example: ${RUNTIME_EXECUTABLE_NAME} --format win7_amd64.iso /dev/sdd"
	echo -e
	echo -e "Options"
	echo -e "\t--verbose, -v        ""Verbose mode"
	echo -e "\t--help, -h           ""Show this help message and exit"
	echo -e "\t--no-color           ""Disable color"
	echo -e "\t--debug, -d          ""Enable script debugging"
}; declare -fr print_help

process_commandline_parameters(){
	util_check_function_parameters_quantity 0 "${#}"

	if [ "${#RUNTIME_COMMANDLINE_PARAMETERS[@]}" -eq 0 ]; then
		return 0
	fi

	local -a parameters=("${RUNTIME_COMMANDLINE_PARAMETERS[@]}")
	local enable_debug=N

	while [ "${#parameters[@]}" -ne 0 ]; do
		case "${parameters[0]}" in
			"--help"\
			|"-h")
				print_help;
				exit 0
				;;
			"--debug"\
			|"-d")
				enable_debug="Y"
				;;
			'--install')
				install_method='edit'
				shift_array parameters
				source_media="${parameters[0]}"
				shift_array parameters
				target_media="${parameters[0]}"
				;;
			'--format')
				install_method='format'
				shift_array parameters
				source_media="${parameters[0]}"
				shift_array parameters
				target_media="${parameters[0]}"
				;;
			'--verbose'|'-v')
				verbose='1'
				;;
			'--forGui')
				no_color='1'
				only_for_gui='1'
				;;
			'--no-color')
				no_color='1'
				;;
			*)
				printf "ERROR: Unknown command-line argument \"%s\"\n" "${parameters[0]}" >&2
				return 1
				;;
		esac
		shift_array parameters
	done

	if [ "${enable_debug}" = "Y" ]; then
		trap 'trap_return "${FUNCNAME[0]}"' RETURN
		set -o xtrace
	fi
	return 0
}; declare -fr process_commandline_parameters;

shift_array(){
	util_check_function_parameters_quantity 1 "${#}"

	local -n array_ref="${1}"

	if [ "${#array_ref[@]}" -eq 0 ]; then
		printf "ERROR: array is empty!\n" 1>&2
		return 1
	fi

	# Unset the 1st element
	unset "array_ref[0]"

	# Repack array if element still available in array
	if [ "${#array_ref[@]}" -ne 0 ]; then
		array_ref=("${array_ref[@]}")
	fi

	return 0
}; declare -fr shift_array

trap_return(){
	util_check_function_parameters_quantity 1 "${#}"

	local returning_function="${1}"

	printf "DEBUG: %s: returning from %s\n" "${FUNCNAME[0]}" "${returning_function}" 1>&2
}; declare -fr trap_return

init "${@}"