#!/bin/bash

#    New_Ikki_USB_Creator_en.sh - Create a bootable USB key with a FAT32 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: 1.2 - 02/15/2024

export LANG=C
set -e
label_fat32="IKKIBOOT"
label_exfat="Data"
size_fat32=4
script=$0
mount_ikki=/media/mount_ikki
mount_usb=/media/mount_usb

# Aide
usage() {
cat <<EOF

Usage: $0 -i [Iso] -L [Label]
Options:
  -i	ISO file name
  -L	Key name, default ${label_exfat}
  -s    Default Ikki Boot partition size in GB ${size_fat32}
  -u    Update
  -t	Destination device name (Expert)
  -D	Debug mode
  -h	Displays this message
	
Examples:
  $0 -i Ikki_Boot.iso
  $0 -i Ikki_Boot.iso -L Job
  $0 -i Ikki_Boot.iso -L Job -s 8
  $0 -i Ikki_Boot.iso -u
  $0 -i Ikki_Boot.iso -t /dev/sdX
  $0 -i Ikki_Boot.iso -L Job -t /dev/sdX
	
EOF
}

# Options management
while getopts ":i:L:s:t:Dhu" option; do
    case "${option}" in
        i)
			iso=${OPTARG}
			;;
		L)
			label_exfat=${OPTARG}
			;;
		s)
			size_fat32=${OPTARG}
			;;
		u)
			update=true
		;;
		t)
			target=${OPTARG}
			;;
		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 whether the -s option has been used and tests the value
if [[ -n "${size_fat32}" ]]; then
	if [[ "${size_fat32}" -eq "${size_fat32}" ]] 2>/dev/null ; then
		if [[ "${size_fat32}" -lt 3 ]] || [[ "${size_fat32}" -gt 20 ]]; then
			echo
			echo "Partition size must be between 3 and 20 GB"
			echo
			exit 1
		fi
	else
		echo
		echo "Partition size must be an integer"
		echo
		exit 1
	fi
fi

# Checks the number of characters in the label
if [ "${#label_fat32}" -gt 11 ]; then 
	echo
	echo "The name of the first partition of the key must not exceed 11 characters." 
	echo
	exit 1
fi

# Test if the file is an .iso
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 doesn't 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
	{
		# List removable disk devices
		list_key=$(lsblk -nbdpo NAME,SIZE,TYPE,HOTPLUG | grep -Ev " 0.disk" | grep -E ".disk.*1$")
	} || {
		echo
		echo "No USB key detected"
		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 "A single 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 "Destination device ${target} 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 - if so, dismantles it
target_use=$(mount | grep "${target}" | awk '{print $3}')
if [[ -n "${target_use}" ]]; then
	echo
	echo "The USB key is currently in use and will be unmounted."
	set +e
	for use in ${target_use}; do umount "${use}";done
	set -e
fi

# If the key is still in use, we quit
target_use=$(mount | grep "${target}" | awk '{print $3}')
if [[ -n "${target_use}" ]]; then
	echo
	echo "The USB key is always in use :"
	echo "${target_use}"
	echo
	echo "Quit all applications that use this key and try again."
	echo
	exit 1
fi

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

# Checks key size against 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

# Copy the contents of the ISO image to the USB key
target_space=$((target_space/1024/1024/1024))

while [[ "${confirmation}" != "YES" ]] && [[ "${confirmation}" != "NO" ]] ; do
	echo
	if [[ "${update}" = true ]]; then
		read -rp "Are you sure you want to update the ${target} disk from ${target_space} Go ? (YES/NO) : " confirmation
	else
		read -rp "Are you sure you want to delete the ${target} disk from ${target_space} Go ? (YES/NO) : " confirmation
	fi
done
if [[ "${confirmation}" == "YES" ]] ; then
	echo
	
        if [[ "${update}" = true ]]; then
		if [[ ! -d "${mount_ikki}" ]]; then
                	mkdir -p "${mount_ikki}"
        	        to_delete_iso=true
	        fi
        	mount -o loop "${iso}" "${mount_ikki}" >/dev/null 2>&1

        	if [[ ! -d "${mount_usb}" ]]; then
                	mkdir -p "${mount_usb}"
                	to_delete_usb=true
        	fi
        	if ! mount -t vfat "${target}1" "${mount_usb}" >/dev/null 2>&1; then
			echo "This USB key was not created with the script $0"
			echo "Try again without the -u option"
			echo
			sleep 2
			umount "${mount_ikki}" >/dev/null 2>&1
                        umount "${mount_usb}" >/dev/null 2>&1
			exit 1
		fi
                
		if [[ ! -f "${mount_usb}/doc/licence/md5sum.txt" ]]; then
			echo "This USB key does not contain Ikki Boot, update canceled"
			echo
			sleep 2
			umount "${mount_ikki}" >/dev/null 2>&1
			umount "${mount_usb}" >/dev/null 2>&1
			exit 1
		fi

		# Removal of the previous version while retaining any Porteus modules that may have been added
		echo "1) Remove previous version"
		echo
		echo "Removing Porteus modules"
		
		# Lists and deletes official xzm files
		list_to_delete=$(cat "${mount_usb}/doc/licence/md5sum.txt" | grep -E ".*xzm$" | awk '{print $2}' | cut -c2-)
                for file in $list_to_delete
		do
  			if [[ -f "${mount_usb}${file}" ]]; then
				rm "${mount_usb}${file}"
			fi
		done

		# Delete Porteus files and directories except base, modules and optional
		list_to_delete=$(cat "${mount_usb}/doc/licence/md5sum.txt" | awk '{print $2}' | grep -E "^./porteus" | cut -d '/' -f3 | uniq | grep -Ev "base|modules|optional")
                for file in $list_to_delete
                do
			rm -rf "${mount_usb}/porteus/${file}"
                done

		# Deleting other directories
		 list_to_delete=$(cat "${mount_usb}/doc/licence/md5sum.txt" | cut -d '/' -f2 | uniq | grep -Ev "porteus|extra")
                for file in $list_to_delete
                do
                        rm -rf "${mount_usb}${file}"
			echo "Suppression $(basename "${file}")"
			sync
                done

		echo
		echo "2) Copy ISO image content to key"
	        echo
 
		# Shows copy progress
        	find "${mount_ikki}" -mindepth 1 -maxdepth 1 \! -type l -exec bash -c 'echo Copy $(basename "$1")' shell {} \; -exec cp -rf "{}" "${mount_usb}" \; -exec sync \;
        	echo
        	echo "Copy completed"

		sync

		echo
                echo "3) Install Syslinux"
                cp -a "${mount_usb}"/script/syslinux/linux /tmp
                cd /tmp/linux

                # Install Syslinux
                chmod +x extlinux
                ./extlinux -i "${mount_usb}"/boot/syslinux/

                # Copy the Syslinux MBR
                dd conv=notrunc bs=440 count=1 if=mbr.bin of="${target}"

                # Add boot flag
                sfdisk -A "${target}" 1 >/dev/null 2>&1

        	umount "${mount_ikki}"
	        umount "${mount_usb}"

        	# Cleaning
        	if [[ "${to_delete_usb}" = true ]]; then
                	rm -rf "${mount_usb}"
        	fi

        	if [[ "${to_delete_iso}" = true ]]; then
                	rm -rf "${mount_ikki}"
        	fi

        	sync

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

		exit 0
        fi

	# Creating partitions
	echo "1) Create partitions"
	(
	echo o                  # create a new empty MBR (DOS) partition table
	echo n                  # Add a new partition
	echo p                  # Primary partition
	echo 1			        # Partition number
	echo                    # First sector (Accept default)
	echo +"${size_fat32}"G  # Last sector (4 Go)
	echo t 			        # Type
	echo 0c                 # FAT32
	echo n                  # Add a new partition
	echo p                  # Primary partition
	echo			        # Partition number
	echo                    # First sector (Accept default)
	echo                    # Last sector (Accept default)
	echo t                  # Type
	echo 2                  # Second partition
	echo 07                 # ExFAT
	sleep 5 || true
	echo w                  # Write changes
	echo q                  # Quit
	) | fdisk -w always "${target}"
	
	until fdisk -l | grep "${target}"1 1> /dev/null
	do
	  sleep 1
	done
	
	# Format 1st partition as FAT32
	label_fat32=$(echo "${label_fat32}" | tr "[:lower:]" "[:upper:]")
	mkfs.vfat -F 32 -n "${label_fat32}" "${target}1"
	
	# Format 2nd partition as exFAT
	mkfs.exfat -L "${label_exfat}" "${target}2"
	
    echo
	echo "2) Copy ISO image content to key"
	echo
	if [[ ! -d "${mount_ikki}" ]]; then
		mkdir -p "${mount_ikki}"
		to_delete_iso=true
	fi
	mount -o loop "${iso}" "${mount_ikki}" >/dev/null 2>&1
	
	if [[ ! -d "${mount_usb}" ]]; then
		mkdir -p "${mount_usb}"
		to_delete_usb=true
	fi
	mount -t vfat "${target}1" "${mount_usb}"
	
	# Displays copy progress
	#for file in "${mount_ikki}"/*; do echo Copie "${file}"; cp -rfL "${file}" "${mount_usb}";sync;done
	find "${mount_ikki}" -mindepth 1 -maxdepth 1 \! -type l -exec bash -c 'echo Copy $(basename "$1")' shell {} \; -exec cp -rf "{}" "${mount_usb}" \; -exec sync \;
	echo
	echo "Copy completed"
	
	umount "${mount_ikki}"
	
	echo
	echo "3) Install Syslinux"
	cp -a "${mount_usb}"/script/syslinux/linux /tmp
	cd /tmp/linux
	
	# Install Syslinux
	chmod +x extlinux
	./extlinux -i "${mount_usb}"/boot/syslinux/
	
	# Copy the Syslinux MBR
	dd conv=notrunc bs=440 count=1 if=mbr.bin of="${target}"
	
	# Add boot flag
	sfdisk -A "${target}" 1 >/dev/null 2>&1
	
	sync
	umount "${mount_usb}"
	
	# Cleaning
	if [[ "${to_delete_usb}" = true ]]; then
		rm -rf "${mount_usb}"
	fi
	
	if [[ "${to_delete_iso}" = true ]]; then
		rm -rf "${mount_ikki}"
	fi
	
	rm -rf /tmp/linux
	
else
	exit 1
fi

sync

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