#!/bin/bash

#    Ikki_USB_Creator_en.sh - Create a bootable USB key with a persistent partition from an iso file
#    Copyright (C) 2007-2024  Julien Billard

#    This program 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.

#    This program 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 this program.  If not, see <https://www.gnu.org/licenses/>.

# https://ikkiboot.tuxfamily.org
# Version: 2.7 - 06/29/2023

export LANG=C
set -e
label="Data"
script=$0
block_size=4

# Help
usage() {
cat <<EOF

Usage: $0 -i [Iso] -L [Label]                                                                
Options:                                                                                     
  -i    ISO file name
  -L    Name of the persistent partition, by default "${label}"
  -t    Name of the destination device (Expert)
  -b    Block size in M for dd, default "${block_size}" (Expert)
  -D    Mode debug                                                              
  -h    Displays this message  

Examples:
  $0 -i Ikki_Boot.iso
  $0 -i Ikki_Boot.iso -L Job
  $0 -i Ikki_Boot.iso -b 10
  $0 -i Ikki_Boot.iso -t /dev/sdX
  $0 -i Ikki_Boot.iso -L Job -b 10
  $0 -i Ikki_Boot.iso -L Job -t /dev/sdX -b 10
	
EOF
}

# Options management
while getopts ":i:L:t:b:Dh" option; do
    case "${option}" in
        i)
			iso=${OPTARG}
			;;
		L)
			label=${OPTARG}
			;;
		t)
			target=${OPTARG}
			;;
		b)
			block_size=${OPTARG}
			;;
	        D)
			set -x
			;;
		h|*)
			usage
			exit 1
			;;
    esac
done

# Checks if the user is root
if [[ "${UID}" -ne 0 ]]; then
	echo
	echo "[${LOGNAME}] You must run this script as root."
	echo
	exit 1
fi

# Checks if the fdisk command works
if ! [[ -x "$(command -v fdisk)" ]]; then
        echo
        echo "You must log in as root with the command :"
        echo "su -"
        echo
        exit 1
fi

# Checks if a parameter has been used
if [[ $# -eq 0 ]]; then 
	usage
	exit 1
fi

# Checks if the -i option has been used
if [[ -z "${iso}" ]]; then
	usage
	exit 1
fi

# Checks if the -b option has been used and tests the value
if [[ -n "${block_size}" ]]; then
	if [[ "${block_size}" -eq "${block_size}" ]] 2>/dev/null ; then
		if [[ "${block_size}" -lt 1 ]] || [[ "${block_size}" -gt 100 ]]; then
			echo
			echo "The size of the blocks must be between 1 and 100"
			echo
			exit 1
		fi
	else
		echo
		echo "The size of the blocks must be an interger"
		echo
		exit 1
	fi
fi

# Test if the file is an.iso file
len=${#iso}
len=$((len-3))
if [[ "${iso:len}" != "iso" ]]; then
	usage
	exit 1
fi

# Checks if the ISO file exists
if [[ ! -f "${iso}" ]]; then
	echo
    echo "The ${iso} file does not exist";
	echo
    exit 1
fi

# Checks if the ISO file is in the same directory as the script
if [[ $(dirname "${iso}") != $(dirname "${script}") ]]; then
	echo
	echo "The ISO file must be in the same directory as this script"
	echo
    exit 1
fi

if [[ -z "${target}" ]]; then
	{
		# Lists removable disk devices
		list_key=$(lsblk -nbdpo NAME,SIZE,TYPE,HOTPLUG | grep -Ev " 0.disk" | grep -E ".disk.*1$")
	} || {
		echo
		echo "No USB flash drives have been found"
		echo
		exit 1
	}

	# Number of removable disks detected
	nbr_word=$(echo "${list_key}" | wc -w)
	nbr_key=$((nbr_word / 4))

	if [[ "${nbr_key}" -eq 1 ]]; then
		target=$(echo "${list_key}" | awk '{print $1}')
		target_space=$(echo "${list_key}" | awk '{print $2}')
		echo
		echo "A USB key has been detected :"

		lsblk -dpo MODEL,NAME,SIZE | grep -E "${target}"
		echo
		while [[ "${confirmation}" != "y" ]] && [[ "${confirmation}" != "n" ]] ; do
			read -rp "Do you confirm the use of this key y/n ? " confirmation
		done
	elif [[ "${nbr_key}" -gt 1 ]]; then
		echo
		echo "Several USB keys have been detected :"
		echo
		lsblk -dpo MODEL,NAME,SIZE,TYPE,HOTPLUG | grep -Ev " 0.?.disk" | grep -E ".disk.*1$" | rev | cut -d ' ' -f3- | rev | sed s/disk//g
		echo
		echo "Only one USB key must be connected"
		echo
		exit 1
	fi
	
else
	#Test $target
	{
		list_key=$(lsblk -nbdpo NAME,SIZE,TYPE,HOTPLUG | grep -E "${target}.*disk.*1$")
	} || {
		echo
		echo "The ${target} destination device is not removable"
		echo
		exit 1
	}
	echo
	echo "Destination device :"
	lsblk -dpo MODEL,NAME,SIZE | grep -E "${target}"
	target_space=$(echo "${list_key}" | awk '{print $2}')
	echo
	while [[ "${confirmation}" != "y" ]] && [[ "${confirmation}" != "n" ]] ; do
		read -rp "Do you confirm the use of this key y/n ? " confirmation
	done
	
fi

if [[ "${confirmation}" = "n" ]]; then
	exit 1
fi

# Checks if the key is in use
target_use=$(mount | grep "${target}" | awk '{print $3}')
if [[ -n "${target_use}" ]]; then
	echo
	echo "The USB key is currently used in :"
	echo "${target_use}"
	echo
	echo "Unmount the USB key and run the script again"
	exit 1
fi

iso_space=$(du -b "${iso}" | awk '{print $1}')

# Checks the key size in comparison to the ISO file
if [[ "${iso_space}" -gt "${target_space}" ]]; then
	echo
	echo "A larger capacity key is required for this ISO file"
	echo
	exit 1
fi

# We write the ISO image on the USB stick
target_space=$((target_space/1024/1024/1024))

while [[ "${confirmation}" != "YES" ]] && [[ "${confirmation}" != "NO" ]] ; do
	echo
	read -rp "Are you sure you want to delete the ${target} disk of ${target_space} Go (YES/NO) : " confirmation
done
if [[ "${confirmation}" == "YES" ]] ; then
	echo
	echo "1) Writing the ISO image on the key"
	echo
	echo "dd if=${iso} of=${target} bs=${block_size}M status=progress && sync" 
	dd if="${iso}" of="${target}" bs="${block_size}"M status=progress && sync
else
	exit 1
fi
sleep 5
echo
echo "2) Creating an exFAT partition on ${target}"

# Partition number test
count_part=$(lsblk "${target}" | wc -l)
list_count_part=$((count_part - 2))

if [[ "${list_count_part}" -eq 3 ]] ; then
	# Delete the third partition
	(
	echo d                  # Del partition
	echo 3                  # Partition number
	echo w                  # Write changes
	) | fdisk "${target}"
fi
	# Create exFAT partition
	(
	echo n                  # Add a new partition
	echo p                  # Primary partition
	echo 3			# Partition number
	echo                    # First sector (Accept default)
	echo                    # Last sector (Accept default)
	echo t 			# Type
	echo 3		 	# Partition number
	echo 7 			# exFAT
	sleep 5 || true
	echo w                  # Write changes
	echo q                  # Quit
	) | fdisk "${target}"

# Wait for the 3rd partition to be detected
until fdisk -l | grep "${target}"3 1> /dev/null
do
  sleep 1
done

# Format partition in exFAT
mkfs.exfat -L "${label}" -c 4096 "${target}"3
sync

echo
echo "The USB key has been successfully created"
echo