#! /bin/bash

set -e

PARTED="/sbin/parted"
PARTX="/usr/bin/partx"
UEFI_HELPER="/usr/share/ubuntu/scripts/uefi_helper.sh"
SECURE_BOOT_VAR="/sys/firmware/efi/vars/SecureBoot-8be4df61-93ca-11d2-aa0d-00e098032b8c/data"

# We just use the first EFI partition for simple.
EFI_PART=`blkid | grep 'EFI' | head -n1`
EFI_SIZE=`parted -l | grep 'EFI' | head -n1 | awk '{print $3 $4}'`

if [ -z "$EFI_PART" ]; then
	# No efi partition, we don't do anything.
	exit 0
fi
# We just check the updated sutton EFI partition, when it use the 50.0MB.
if [ "$EFI_SIZE" = "50.0MB18.0MB" ]; then

	# Get the EFI partition's name and disk's name
	EFI_PART_NAME=`echo ${EFI_PART} | cut -d ':' -f1`
	if [ "${EFI_PART_NAME:5:2}" = "nv" ]; then
		DISK_NAME=${EFI_PART_NAME:0:-2}
	elif [ "${EFI_PART_NAME:5:2}" = "sd" ]; then
		DISK_NAME=${EFI_PART_NAME:0:-1}
	else
		DISK_NAME=""
	fi

	if [ -z $DISK_NAME ]; then
		exit 0
	fi
	# Backup the /etc/fstab
	cp /etc/fstab /etc/fstab.new.efi
	# Backup EFI partition UUID and PARTUUID
	ORIGIN_EFI_UUID=`echo ${EFI_PART} | grep -o ' UUID="[^ "]\+' | cut -d '"' -f2`

	# Make new efi directory, and backup the EFI files
	EP=/boot/efi.bak
	if ! [ -d $EP ]; then
		mkdir $EP
	fi
	rsync -av /boot/efi/ $EP/

	umount ${EFI_PART_NAME}

	# Remove the EFI partition
	$PARTED -s $DISK_NAME rm 1
	$PARTX -d --nr 1 $DISK_NAME || true

	# Create the new EFI partition with 1M offset
	$PARTED -s $DISK_NAME -a minimal mkpart primary fat32 1M 50M
	$PARTED -s $DISK_NAME name 1 EFI
	# Wait for the parted event queue to finish.
	/bin/udevadm settle
	/sbin/mkfs.msdos -n "EFI System" $EFI_PART_NAME
	$PARTED -s $DISK_NAME set 1 boot on

	# Get the new UUID
	NEW_EFI_UUID=`blkid | grep "${EFI_PART_NAME}" | grep -o ' UUID="[^ "]\+' | cut -d '"' -f2`
	# Update the /etc/fstab
	/bin/sed -i "s/${ORIGIN_EFI_UUID}/${NEW_EFI_UUID}/g" /etc/fstab

	mount ${EFI_PART_NAME} /boot/efi
	# Restore the EFI partition's file
	rsync -av $EP/ /boot/efi/

	# Update the efibootmgr with new PARTUUID
	LABEL="ubuntu"
	LOADER="\\EFI\\ubuntu\\grubx64.efi"
	if [ -e $SECURE_BOOT_VAR ]; then
		LOADER="\\EFI\\ubuntu\\shimx64.efi"
	fi
	sh $UEFI_HELPER add $DISK_NAME 1 $LOADER $LABEL
fi
