#!/bin/sh

set -e

USAGE=$(cat <<HELP
 usage: $0 [-h|--help] [--yes-to-all] fs_type
 
 -h, --help    display help message
 --yes-to-all  skip confirmation
 fs_type       filesystem type of /mnt/flash
               valid options are "ext4" or "vfat"
HELP
)

DESCRIPTION=$(cat <<HELP
 Repartitions the flash device into /mnt/flash, a recovery partition, and if
 the system's filesystem is ext4, /mnt/crash as well. The partition sizes are
 determined by the capacity of the flash device. If fs_type is ext4, an ext4
 filesystem is created on /mnt/flash and /mnt/crash. Otherwise, if fs_type is
 VFAT, a VFAT filesystem is created on /mnt/flash. After running this script,
 an EOS.swi, boot-config, and optional startup-config should be added to
 /mnt/flash followed by rebooting the switch.
HELP
)

# Whether to skip user confirmation
yes_to_all=false

# Filesystem type of /mnt/flash
fs_type=

make_ext4_filesystem() {
   local mount_name=$1
   local mount_point=/mnt/$mount_name
   local label="eos_$mount_name"
   local dev=$2
   local part=$3
   local extra_opts=$4

   echo "Creating an ext4 filesystem on $mount_point..."
   mke2fs -t ext4 -Fq -m 0 -L $label $extra_opts "${dev}${part}" > /dev/null 2>&1

   mount "${dev}${part}" $mount_point
   chown -R root:88 $mount_point
   chmod -R ug=rwX,o= $mount_point
   tune2fs -c 1 "${dev}${part}" > /dev/null 2>&1 || :
   umount $mount_point
}

make_vfat_filesystem() {
   local mount_name=$1
   local mount_point=/mnt/$mount_name
   local label="eos_$mount_name"
   local dev=$2
   local part=$3
   local extra_opts=$4

   echo "Creating a VFAT filesystem on ${1}..."
   mkfs.vfat -F 32 -n $label $extra_opts "${dev}${part}" > /dev/null 2>&1
}

while [[ $# -gt 0 ]]; do
   case "$1" in
      -h|--help)
         echo
         echo "$DESCRIPTION"
         echo
         echo "$USAGE"
         echo
         exit 0
         ;;
      --yes-to-all)
         yes_to_all=true
         shift
         ;;
      *)
         if [ "$1" == "ext4" ] || [ "$1" = "vfat" ]; then
            fs_type="$1"
            shift
         else
            echo "Invalid filesystem type specified"
            exit 1
         fi
         ;;
   esac
done

if [ -z "$fs_type" ]; then
   echo "No filesystem type specified"
   exit 1
fi

# locate flash device to reformat (i.e. /dev/sda)
sys_blocks=$(ls -1 /sys/block)
old_ifs=$IFS
IFS='
'
block_devs=$(cat /etc/blockdev)
for dev in $block_devs; do
   IFS=$old_ifs
   set -- $dev
   dev_re=$1
   if [ "$2" = "flash" ]; then
      for block in $sys_blocks; do
         path=$(realpath /sys/block/$block/device 2> /dev/null || :)
         if echo $path | grep -E "/sys/devices/.*"; then
            dev_id=$(echo $path | sed -r "s/^\/sys\/devices\///")
            if echo $dev_id | grep -E $dev_re; then
               flash_block=$block
               break 2
            fi
         fi
      done
   fi
done
if [ -z "$flash_block" ]; then
   echo "Could not locate flash device...exiting"
   exit 0
fi
flash_dev="/dev/$flash_block"

# get user confirmation to reformat their flash
if ! $yes_to_all; then
   echo "WARNING: Proceeding will reformat all partitions on $flash_dev"
   printf 'To continue, type "yes" and press Enter or just press Enter to cancel: '
   read arg
   if [ "$arg" != "yes" ]; then
      echo "Cancelled"
      exit 1
   fi
fi

# unmount existing partitions that are mounted in the filesystem
echo "Unmounting all of ${flash_dev}'s partitions..."
umount ${flash_dev}* 2> /dev/null || :

# determine usable flash size
SECTORS_PER_KIB=2
SECTORS_PER_MIB=$((SECTORS_PER_KIB * 1024))
disk_info=$(fdisk -l /dev/$flash_block | grep "Disk /dev/$flash_block")
size_bytes=$(echo $disk_info | grep -Eo "[0-9]+ bytes" | sed -r "s/ bytes//")
size_gb=$((size_bytes / 1000000000))
echo "$flash_dev has size $size_gb GB..."
total_sectors=$((size_bytes / 1024 * SECTORS_PER_KIB))
start_sectors=$SECTORS_PER_MIB
crash_sectors=$((64 * SECTORS_PER_MIB))
recovery_sectors=$SECTORS_PER_MIB
main_sectors=$((total_sectors - start_sectors - recovery_sectors))
if [ "$fs_type" == "ext4" ]; then
   main_sectors=$((main_sectors - crash_sectors))
fi

# keep partition layout in ascending order
for row in '35 28560' '130 114400' '260 228870'; do
   set -- $row 
   if [ $size_gb -lt $1 ]; then
      actual_size_mib=$2
      break
   fi
done

# determine the final size of /mnt/flash
product_sectors=$((actual_size_mib * SECTORS_PER_MIB))
if [ $product_sectors -lt $main_sectors ]; then
   main_sectors=$product_sectors
fi

echo "Reserving $main_sectors sectors for /mnt/flash..."
if [ "$fs_type" == "ext4" ]; then
   echo "Reserving $crash_sectors sectors for /mnt/crash..."
fi
echo "Reserving $recovery_sectors sectors for the recovery partition..."

# zero out first MB of flash
dd if=/dev/zero of=$flash_dev bs=512 count=$SECTORS_PER_MIB > /dev/null 2>&1

# define partition ranges
# if fs_type is ext4, we'll need three partitions. One for /mnt/flash,
# /mnt/flash, and recovery
# if fs_type is VFAT, we'll need two partitions. One for /mnt/flash and one
# for recovery
start1=$start_sectors
end1=$((start1 + main_sectors - 1))
fdiskCmdStr="n\np\n1\n$start1\n$end1\na\n1\n" #p1 fdisk cmds
start2=$((end1 + 1))

if [ "$fs_type" == "ext4" ]; then
   fdiskCmdStr=""$fdiskCmdStr"t\n83\n" # Set system id code of p1
   end2=$((start2 + crash_sectors - 1))
   fdiskCmdStr=""$fdiskCmdStr"n\np\n2\n$start2\n$end2\nt\n2\n83\n" # p2 fdisk cmds
   start3=$((end2 + 1))
   end3=$((start3 + recovery_sectors - 1))
   fdiskCmdStr=""$fdiskCmdStr"n\np\n3\n$start3\n$end3\nt\n3\n12\n" # p3 fdisk cmds
else
   # fs_type is VFAT
   fdiskCmdStr=""$fdiskCmdStr"t\nc\n" # Set system id code of p1
   end2=$((start2 + recovery_sectors - 1))
   fdiskCmdStr=""$fdiskCmdStr"n\np\n2\n$start2\n$end2\nt\n2\n12\n" # p2 fdisk cmds
fi

echo "Creating new partitions on ${flash_dev}..."
# Disable auto mounting
>/proc/sys/kernel/hotplug
printf "o\n"$fdiskCmdStr"w\n" | fdisk -u $flash_dev > /dev/null 2>&1

mdev -s
umount ${flash_dev}* 2> /dev/null || :

tmp_recov="/dev/recoveryPartition"
tmp_cpio="/tmp/recovery.cpio"
recover_part=

# eMMC partition names are prefixed with a "p"
part_prefix=
if [ -n "$( echo "$flash_block" | grep "mmc" )" ]; then
   part_prefix="p"
fi
if [ "$fs_type" == "ext4" ]; then
   make_ext4_filesystem "flash" "$flash_dev" "${part_prefix}1"
   make_ext4_filesystem "crash" "$flash_dev" "${part_prefix}2" "-J size=1"
   recover_part="${part_prefix}3"
else
   make_vfat_filesystem "flash" "$flash_dev" "${part_prefix}1"
   recover_part="${part_prefix}2"
fi

# create an empty cpio image for the recovery partition
echo "" | cpio -o -H newc > $tmp_cpio
cat $tmp_cpio > "${flash_dev}${recover_part}"
rm -f $tmp_cpio

if [ "$fs_type" == "ext4" ]; then
   echo "Mounting /mnt/crash and /mnt/flash..."
else
   echo "Mounting /mnt/flash..."
fi
mdev -s
# Re-enable auto mounting
echo /bin/mdev > /proc/sys/kernel/hotplug
echo "Done!"
