#!/bin/bash

if [ $(id -u) != 0 ]; then
    echo "Drive partitioning requires root privileges!"
    exit 1
fi

# this script will write the partition type, selected partition, and selected
# drive into the specified file
OUTFILE=$1

if [ ! -f "$OUTFILE" ]; then
  echo "Output file does not exist. Exiting..."
  exit 1
fi

# source in the functions
source /opt/vyatta/sbin/install-functions

# the INSTALL_LOG env var should be exported by the "caller".
# it will be used to log messages.

# Absolute minimum root partition size in MB.  Below this, we won't let
# you install.
ROOT_MIN=2000
# the base install drive e.g. sda
INSTALL_DRIVE=''
# the install partition e.g. sda1
ROOT_PARTITION=''
# the type of the install partition: "union", "old", or "new"
ROOT_PARTITION_TYPE=''
# global holding variable used in the select_partition sub
PARTITION=''
# default file system type
ROOT_FSTYPE='ext4'

warn_of_dire_consequences () {
  # Give the user a requisite warning that we are about to nuke their drive
  response=''
  while [ -z "$response" ]; do
    echo "This will destroy all data on /dev/$INSTALL_DRIVE."
    echo -n "Continue? (Yes/No) [No]: "
    response=$(get_response "No" "Yes No Y N")

    if [ "$response" == "no" ] || [ "$response" == "n" ]; then
      echo "Ok then.  Exiting..."
      exit 1
    fi
  done
}

check_for_old_raid () {
  # First, trigger construction of previously configured RAID groups
  echo -n "Looking for pre-existing RAID groups..."
  raid_config=`mdadm --examine --scan`
  if [ -z "$raid_config" ]; then
    echo "none found."
    return
  fi

  echo "found some."
  echo "Trying to configure pre-existing RAID groups..."
  mdadm --assemble --scan --auto=yes --symlink=no

  # Identify physical drives
  raid_drives=$(cat /proc/partitions | awk '{ if ($4!="name") { print $4 } }' \
                | grep "md" | egrep -v "^$")
  if [ -z "$raid_drives" ]; then
    echo "Unable to configure any RAID groups."
    return
  fi

  numraids=`echo $raid_drives | wc -w`
  if [ $numraids -eq 1 ]; then
    echo "The following RAID group is now configured:"
  else
    echo "The following RAID groups are now configured:"
  fi

  for drive in $raid_drives; do
    cat /proc/mdstat | grep --after-context 2 ^$drive | sed -e 's/^/\t/'
  done

  if [ $numraids -eq 1 ]; then
    echo -n "Would you like to use this one? (Yes/No) [Yes]:"
  else
    echo -n "Would you like to use one of these? (Yes/No) [Yes]:"
  fi

  response=$(get_response "Yes" "Yes No Y N")
  if [ "$response" == "no" ] || [ "$response" == "n" ]; then
    echo
    echo "Ok.  Not using existing RAID groups."
    echo

    # pick the first RAID group to be broken
    raid_drive=$(echo $raid_drives | /usr/bin/awk '{ print $1 }')

    echo "Would you like to break RAID group $raid_drive so that its"
    echo "members can be re-used for a new installation, understanding"
    echo -n "that doing so will destroy all data on it? (Yes/No) [No]:"
    destroy_raid=$(get_response "No" "Yes No Y N")
    echo

    if [ "${destroy_raid:0:1}" = "y" ]; then
	    echo "OK.  Breaking the RAID group $raid_drive."

	    members=`ls /sys/block/$raid_drive/slaves`

	    echo "First, stopping all existing RAID groups:"
	    mdadm --stop --scan

	    for member in $members ; do
        drive=${member:0:3}
        part=${member:3:1}
        echo "Re-setting partition ID for RAID group $raid_drive member /dev/${member}:"
        sfdisk --change-id /dev/$drive $part 0x83
        echo "Clearing RAID superblock from RAID group $raid_drive member /dev/${member}."
        mdadm --zero-superblock /dev/$member
	    done
    else
	    echo "OK.  Stopping, but not breaking, existing RAID groups:"
	    mdadm --stop --scan
    fi

    echo
    return
  fi

  if [ $numraids -eq 1 ]; then
    INSTALL_DRIVE=$raid_drives
  else
    # take the first drive as the default
    INSTALL_DRIVE=$(echo $raid_drives | /usr/bin/awk '{ print $1 }')
    echo -n "Which one would you like to use? ($raid_drives) [$INSTALL_DRIVE]: "
    INSTALL_DRIVE=$(get_response "$INSTALL_DRIVE" "$drives")
  fi

  echo "Using RAID partition $INSTALL_DRIVE"

  raid_degraded=`cat /sys/block/$INSTALL_DRIVE/md/degraded`
  raid_sync_action=`cat /sys/block/$INSTALL_DRIVE/md/sync_action`

  if [ "$raid_degraded" = "1" ]; then
    echo
    if [ "$raid_sync_action" = "recover" ]; then
	    echo "Error: This RAID set is degraded and is in the process of"
	    echo "rebuilding.  It is not safe to install onto it while the"
	    echo "rebuild is in progress. Please wait for the rebuild to"
	    echo "complete and then re-start the installation.  You may"
	    echo "monitor the progress of the RAID rebuild with the"
	    echo "command:"
	    echo
	    echo "    show raid $INSTALL_DRIVE"
	    echo
	    exit 1
    fi
    echo "Warning:  This RAID set is degraded, but is not in the"
    echo "process of rebuilding.  It is safe to perform the installation"
    echo "onto a degraded RAID set that is not in the process of"
    echo "rebuilding.  You may stop the installation now and rebuild the"
    echo "RAID set, or continue installing onto it.  If you continue"
    echo "installing, do not attempt to rebuild the RAID set until the"
    echo "installation has completed and you have rebooted the system."
    echo
  fi

  warn_of_dire_consequences

  ROOT_PARTITION=$INSTALL_DRIVE
  ROOT_PARTITION_TYPE=new

  # make sure we aren't working on a mounted part
  unmount "$INSTALL_DRIVE"

  # check for an old config on the partition
  mkdir -p /mnt/tmp
  output=$(mount -r /dev/"$ROOT_PARTITION" /mnt/tmp 2>&1)
  if [ $? != 0 ]; then
      echo "Cannot mount $ROOT_PARTITION to check for old config file"
      lecho "Cannot mount $ROOT_PARTITION"."\n"
      lecho "mount $ROOT_PARTITION /mnt/tmp\n"
      lecho "$output"
  else
      save_old_info "$ROOT_PARTITION"
      umount /mnt/tmp
  fi

  # create the filesystem on the part
  make_filesystem "$ROOT_PARTITION"
}

check_for_new_raid () {
  # Identify physical drives
  drives=$(cat /proc/partitions | awk '{ if ($4!="name") { print $4 } }' \
           | egrep -v "[0-9]$" | egrep -v "^$")

  for instdrv in $drives
  do
      if mount | grep iso9660 | grep -q $instdrv
      then
          drives=${drives//"$instdrv"/}
      fi
  done

  numdrives=`echo $drives | wc -w`

  # Need at least two drives for RAID-1.  We don't yet have the code
  # to handle selection of two from a set of 3 or more, so for now, we
  # only support two drives.
  #
  if [ $numdrives -ne 2 ]; then
    return
  fi

  drive1=`echo $drives | awk '{ print $1 }'`
  drive2=`echo $drives | awk '{ print $2 }'`

  drivesize1=$(get_drive_size $drive1)
  drivesize2=$(get_drive_size $drive2)

  # Both drives must have enough space to hold our minimum root filesystem
  #
  if [ $drivesize1 -lt $ROOT_MIN -o $drivesize2 -lt $ROOT_MIN ]; then
    return
  fi

  echo "You have two disk drives:"
  echo -e "\t$drive1 \t$drivesize1 MB"
  echo -e "\t$drive2 \t$drivesize2 MB"

  echo -n "Would you like to configure RAID-1 mirroring on them? (Yes/No) [Yes]:"
  response=$(get_response "Yes" "Yes No Y N")
  if [ "$response" == "no" ] || [ "$response" == "n" ]; then
    echo "Ok.  Not configuring RAID-1."
    return
  fi

  if [ $drivesize1 -ne $drivesize2 ]; then
    echo "Since the disks are not the same size, we will use the smaller"
    echo "of the two sizes in configuring the RAID-1 set.  This will"
    echo "waste some space on the larger drive."
    echo ""
  fi

  # Configure RAID-1
  echo "This process will erase all data on both drives."
  echo -n "Are you sure you want to do this? (Yes/No) [No]: "
  response=$(get_response "Yes" "Yes No Y N")
  if [ "$response" == "no" ] || [ "$response" == "n" ]; then
    echo "Ok.  Not configuring RAID-1."
    return
  fi

  for drive in $drives; do
    echo "Deleting old partitions on drive $drive"
    # remove any existing partitions on that drive
    delete_partitions "$drive"
  done

  # Need to leave space on both disks between the MBR and the start
  # of the first partition for grub. Grub needs to embed a large
  # boot image there when booting off RAID devices.
  #
  # Partition creation variables are in units of megabytes.
  part_start_offset=2
  data_dev=1

  if [ $drivesize1 -lt $drivesize2 ]; then
    root_size=$drivesize1
  else
    root_size=$drivesize2
  fi

  let root_size-=$part_start_offset

  for drive in $drives; do
    create_partitions "$drive" $root_size "no"
    if [ -d /sys/firmware/efi ]; then
        #EFI moves the data parition on RAID to 3
        data_dev=3
        echo "Create data partition: /dev/${drive}${data_dev}"
    else
        echo "Creating data partition: /dev/${drive}${data_dev}"
        sfdisk --change-id /dev/$drive $data_dev 0xfd
        # mark data partition as bootable
        lecho "Marking /dev/$drive partition $data_dev bootable"
        output=$(parted -s /dev/$drive set $data_dev boot on 2>&1)
        lecho "$output"
    fi
  done

  # Must give partition device time to settle
  sleep 5
  echo

  for drive in $drives; do
    echo "Erasing any previous RAID metadata that may exist on /dev/${drive}${data_dev}"
    mdadm --zero-superblock /dev/${drive}${data_dev}
  done

  echo "Creating RAID-1 group on partitions: /dev/${drive1}${data_dev} /dev/${drive2}${data_dev}"

  raid_dev=md0
  yes|mdadm --create /dev/$raid_dev --level=1 --raid-disks=2 --metadata=0.90 \
    /dev/${drive1}${data_dev} /dev/${drive2}${data_dev}

  if [ $? = 0 -a -e /dev/$raid_dev ]; then
    echo "RAID-1 group created successfully:"
    cat /proc/mdstat | grep --after-context 2 ^$raid_dev | sed -e 's/^/\t/'
    if [ -e /usr/sbin/update-initramfs.orig.initramfs-tools ]; then
      echo "Updating initramfs to include the raid config:"
      /usr/sbin/update-initramfs.orig.initramfs-tools -u
    fi
  else
    echo "Unable to create RAID-1 group!"
    return
  fi

  INSTALL_DRIVE=$raid_dev
  ROOT_PARTITION=$INSTALL_DRIVE
  ROOT_PARTITION_TYPE=new

  # Give device time to settle...
  sleep 5

  # create the filesystem on the part
  make_filesystem "$ROOT_PARTITION"
}

# Allow the user to select a partition to work with
# sets the global PARTITION
# $1 is the text to display before prompt
select_partition () {
  minsize=$1
  text=$2
  exclude=$3

  echo -n "Looking for appropriate partitions:  "
  progress_indicator start

  # initialize out global var.  using globals in this way is bad form.  I know.
  PARTITION=''

  # list only the partitions in /proc/partitions.
  parts=$(cat /proc/partitions | awk '{ if ($4!="name") { print $4 " "} }' \
          | egrep "[0-9]" | egrep -v "loop" | tr -d '\n')

  # remove any partitions we have already previously used
  if [ -n "$exclude" ]; then
    for part in $parts; do
      temp=$(echo $part | egrep -v $exclude)
      parts_temp="$parts_temp $temp"
    done
    parts=$parts_temp
  fi

  # Get the partition sizes for display
  # only show linux partitions that have sizes, i.e. remove loops
  display=''
  myparts=''
  for part in $parts; do
    if [ ${part:0:2} = "md" ]; then
      parttype="RAID"
    else
      rootdev=$(echo $part | sed -E 's/p?[0-9]$//g')
      parttype=$(fdisk -l /dev/$rootdev | grep $part | grep Linux)
    fi
    if [ -n "$parttype" ]; then
      lsize=$(get_drive_size $part)
      if [ "$lsize" -a $lsize -ge $minsize ]; then
        display="$display $part\t\t$lsize"MB"\n"
        myparts="$myparts $part"
      fi
    fi
  done

  progress_indicator stop
  echo "OK"

  if [ -n "$myparts" ]; then
    lpartition=''
    while [ -z "$lpartition" ]; do
      # take the first partition as the default
      lpartition=$(echo $myparts | /usr/bin/awk '{ print $1 }')

      echo "I found the following partitions suitable for the VyOS image:"
      echo -e "Partition\tSize"
      echo -e "$display"
      echo
      echo -n "$text [$lpartition]: "

      lpartition=$(get_response "$lpartition" "$myparts")
      echo
    done
  else
    becho "No suitable partition sizes found.  Exiting..."
    exit 1
  fi
  PARTITION=$lpartition
}

rename_old_config() {
  files=$(find /mnt/config -mindepth 1 -type f | grep -v pre-glendale)
  for f in $files; do
    if grep -q '/\*XORP Configuration File, v1.0\*/' $f >&/dev/null; then
      CURTIME=$(date +%F-%H%M%S)
      mv $f $f.pre-glendale.$CURTIME
    fi
  done
}

## save_old_config
# Copy config files from /mnt/tmp to /mnt/config
save_old_config() {
  local part=$1
  local response=''

   # Look to see if there is a config partition there
  while [ -z "$response" ]; do
        echo "$part has an old configuration directory!"
        echo -ne "Would you like me to save the data on it\nbefore I delete it? (Yes/No) [Yes]: "
        response=$(get_response "Yes" "Yes No Y N")
  done

  if [ "$response" == "yes" ] || [ "$response" == "y" ]; then
      mkdir -p /mnt/config
      if [ -d /mnt/tmp/opt/vyatta/etc/config ]; then
          output=$(cp -pR /mnt/tmp/opt/vyatta/etc/config/* /mnt/config)
      else
          output=$(cp -pR /mnt/tmp/* /mnt/config)
      fi
      if [ -n "$output" ]; then
          echo -e "Warning: error in copying the old config partition.\nSee $INSTALL_LOG for more details."
            lecho "Warning: error in copying the old config partition.\ncp -pR /mnt/tmp/* /mnt/config\n$output\n"
      fi
      rename_old_config
  fi
}

## save_old_keys
# Copy SSH keys from /mnt/tmp/etc/ssh /mnt/ssh
save_old_keys() {
    local part=$1
    local response=''

    while [ -z "$response" ]
    do
        echo "$part has SSH host keys"
        echo -ne "Would you like me to keep SSH keys on new install? (Yes/No) [Yes]: "
        response=$(get_response "Yes" "Yes No Y N")
    done

    if [ "$response" == "yes" ] || [ "$response" == "y" ]; then
        mkdir -p /mnt/ssh
	output=$(cp -p /mnt/tmp/etc/ssh/ssh_host_* /mnt/ssh)

        if [ -n "$output" ]; then
            echo -e "Warning: error in copying the old ssh keys."
	    echo -e "See $INSTALL_LOG for more details."
            echo "Warning: error in copying the old ssh keys." >> $INSTALL_LOG
	    echo "cp -pR /mnt/tmp/etc/ssh/ssh_host_* /mnt/ssh" >> $INSTALL_LOG
	    echo "$output\n">> $INSTALL_LOG
	    return
        fi

	# reset modes on keys (should already be set)
	chmod 600 /mnt/ssh/*_key 2>&1
	chmod 644 /mnt/ssh/*.pub 2>&1
	echo "SSH keys have been saved."
    else
	echo "OK.  SSH keys not saved."
    fi
}

save_image_config() {
    image_name=$1

    # Cleanup from possible partial last run
    rm -fr /mnt/config
    mkdir /mnt/config

    output=$(cp -pR /mnt/tmp/boot/$image_name/rw/config/* /mnt/config)

    if [ -n "$output" ]; then
        echo -e "Warning: error in copying the old config partition.\nSee $INSTALL_LOG for more details."
        lecho "Warning: error in copying the old config partition.\ncp -pR /mnt/tmp/* /mnt/config\n$output\n"
    fi
    rename_old_config
}

save_image_keys() {
    image_name=$1

    if [ ! -d /mnt/tmp/boot/$image_name/rw/etc/ssh ]; then
	echo "No SSH keys found on $image_name, so none can be saved."
	return;
    fi

    echo -n "Would you like to save SSH keys from $image_name too? (Yes/No) [Yes] "

    response=$(get_response "Yes" "Yes No Y N")
    if [ "$response" != "yes" ] && [ "$response" != "y" ]; then
	echo "OK. SSH keys not saved."
	return
    fi

    mkdir -p /mnt/ssh
    output=$(cp -p /mnt/tmp/boot/$image_name/rw/etc/ssh/ssh_host_* /mnt/ssh)

    if [ -n "$output" ]; then
        echo -e "Warning: error in copying the old ssh keys."
	echo -e "See $INSTALL_LOG for more details."
        echo "Warning: error in copying the old ssh keys." >> $INSTALL_LOG
	echo "cp -pR /mnt/tmp/etc/ssh/ssh_host_* /mnt/ssh" >> $INSTALL_LOG
	echo "$output\n">> $INSTALL_LOG
	return
    fi

    # reset modes on keys (should already be set)
    chmod 600 /mnt/ssh/*_key 2>&1
    chmod 644 /mnt/ssh/*.pub 2>&1

    echo "SSH keys have been saved."
}

save_old_info() {
    part=$1
    copied=0

    # Cleanup from possible partial last run
    rm -fr /mnt/config
    rm -fr /mnt/ssh

    echo "Looking for config files from previous installations on $part..."

    # Look to see if this is a disk-based installation config partition
    if [ -f /mnt/tmp/opt/vyatta/etc/config/.vyatta_config ] \
	|| [ -f /mnt/tmp/.vyatta_config ]; then
 		save_old_config $1
		copied=1
    fi

    if [ -d /mnt/tmp/etc/ssh ];  then
	save_old_keys $1
	copied=1
    fi

    if [ $copied -eq 0 ]; then
	# Check for images
	images=()
	image_dirs=/mnt/tmp/boot/*

	for dir in $image_dirs; do
	    if [ -f $dir/rw/config/.vyatta_config ]; then

		item=${dir##/mnt/tmp/boot/}
		images=($item ${images[@]})
	    fi
	done

	num_images=${#images[@]}

 	if [ $num_images -gt 0 ]; then
		echo "I found the following installed system image(s) with config files on $part:"
		for (( i = 0; i < $num_images; i++ )); do
			echo "  $((i + 1)): ${images[$i]}"
		done

		if [ $num_images -eq 1 ]; then
		    echo -n "Would you like to save config information from it? (Yes/No) [Yes] "
		else
		    echo -n "Would you like to save config information from one? (Yes/No) [Yes] "
		fi

		response=$(get_response "Yes" "Yes No Y N")
		if [ "$response" != "yes" ] && [ "$response" != "y" ]; then
		    echo "OK.  Config information not saved."
		    return
		fi

		if [ $num_images -gt 1 ]; then
		    image_index=-1
		    while [ $image_index -lt 0 -o $image_index -ge $num_images ];  do
			echo -n "From which image would you like to save config information? (1 .. $num_images): "
			read num
			image_index=$((num - 1))
		    done
		else
		    image_index=0
		fi

		image_name=${images[$image_index]}
		echo "Saving config information from image $image_name."
		save_image_config $image_name
		echo "Done."

		save_image_keys $image_name
	fi
    fi
}


# Delete all existing partitions for an automated install
# $1 is the drive to delete partitions from
delete_partitions () {
  ldrive=$1

  # get the partitions on the drive
  # in the first grep below we add the optional [p] in order to
  # accomdate cciss drives
  partitions=$(cat /proc/partitions | grep $ldrive[p]*[0-9] \
               | awk '{ print $4 }' | sed 's/\(.*\)\([0-9]$\)/\2/g' \
               | grep -v "^$")
  mkdir -p /mnt/tmp

  # now for each part, blow it away
  for lpart in $partitions; do
        dev_name=/dev/$ldrive$lpart
  	output=$(mount -r $dev_name /mnt/tmp 2>&1)
	if [ $? != 0 ]; then
	    echo "cannot mount $dev_name"
	    lecho "Cannot mount $dev_name"."\n"
	    lecho "mount $dev_name /mnt/tmp\n"
	    lecho "$output"
	else
	    save_old_info $ldrive$lpart
	    umount /mnt/tmp
	fi

	lecho "Removing partition $lpart on /dev/$ldrive"
	output=$(parted -s /dev/$ldrive rm $lpart)
	status=$?
	if [ "$status" != 0 ]; then
	    echo -e "Warning: cannot delete partition $lpart on $ldrive.\n"
	    echo -e "Please see $INSTALL_LOG for more details."
	    lecho "Warning: cannot delete partition $lpart on $ldrive.\n"
	    lecho "parted /dev/$ldrive rm $lpart\n$output"
	fi

	# We add a bogus sleep here because the loop needs to wait for udev
	sleep 5
  done
}

# make a filesystem on the drive
# $1 is the drive to format
make_filesystem () {
  ldrive=$1

  echo -n "Creating filesystem on /dev/$ldrive:  "
  lecho "Creating filesystem on /dev/$ldrive..."

  progress_indicator start
  output=$(mkfs -L persistence -t $ROOT_FSTYPE /dev/$ldrive 2>&1)
  status=$?
  if [ "$status" != 0 ]; then
    echo -e "Error: couldn't create the root filesystem.\nSee $INSTALL_LOG for further details.\nExiting..."
    lecho "Error: couldn't create the root filesystem.\n/sbin/mke2fs -j /dev/$ldrive\n$output"
    exit 1
  fi
  progress_indicator stop
  echo "OK"
}

# create the root partition
# $1 is the install drive e.g. sda
# $2 is the partition size e.g. 512
# This will set the global ROOT_PARTITION
create_partitions() {
  ldrive=$1
  root_part_size=$2
  initialize_fs=$3

  # Make sure there is enough space on drive
  size=$(get_drive_size "$ldrive")
  if [ "$root_part_size" -gt "$size" ]; then
    echo "Error: $ldrive is only $size"MB" large.  Desired root is $root_part_size"
    exit 1
  fi
  if [ -d /sys/firmware/efi ]; then
       #Need room for the EFI partition.  512 is standard, but 256 is probably okay here
       root_part_size=$((root_part_size - 256))

       ##Do GPT/EFI Setup
       sgdisk --zap-all /dev/$ldrive
       # part1 = BIOS BOOT (backwards compatibility)
       # part2 = EFI
       # part3 = ROOT
       sgdisk -a1 -n1:34:2047   -t1:EF02 \
           -n2:2048:+256M -t2:EF00 \
           -n3:0:0:+$root_part_size -t3:8300 /dev/$ldrive
       status=$?
       if [ "$status" != 0 ]; then
          echo -e "Error creating primary partition on $ldrive.\nPlease see $INSTALL_LOG for more details.\nExiting..."
          lecho "Error creating primary partition on $ldrive.\nparted /dev/$ldrive mkpart primary 0% $root_part_size\n$output"
          exit 1
       fi
       # set the partition number on the device.
       if [ -n "$( echo $ldrive | grep -E "cciss|ida|nvme|mmcblk" )" ]; then
          # if this is a cciss
          ROOT_PARTITION=$ldrive"p3"
          efipart=$ldrive"p2"
       else
          # else... the rest of the world
          ROOT_PARTITION=$ldrive"3"
          efipart=$ldrive"2"
       fi
       #Add the drive to the file so grub can install
       echo $efipart >> /tmp/efiparts.tmp
   else
      # Force FAT label creation
      lecho "Creating a new disklabel on $ldrive"
      parted -s /dev/$ldrive mklabel msdos

      # Make sure you can print disk info using parted
      parted --script /dev/$ldrive p >/dev/null 2>&1

      # If we still can't, something has gone terribly wrong
      if [ "$?" != "0" ]; then
        echo "Unable to read disk label.  Exiting."
        exit 1
      fi

      lecho "Creating root partition on /dev/$ldrive"

      # Make the root partition
      # if optimal_io_size is empty use default of 2048s
      if [ $(cat /sys/block/$ldrive/queue/optimal_io_size) -gt 0 ]; then
        output=$(parted --script --align optimal /dev/$ldrive mkpart primary 0% $root_part_size)
      else
        output=$(parted --script --align optimal /dev/$ldrive mkpart primary 2048s $root_part_size)
      fi
      status=$?
      if [ "$status" != 0 ]; then
        echo -e "Error creating primary partition on $ldrive.\nPlease see $INSTALL_LOG for more details.\nExiting..."
        lecho "Error creating primary partition on $ldrive.\nparted /dev/$ldrive mkpart primary 0% $root_part_size\n$output"
        exit 1
      fi

      # set the partition number on the device.
      if [ -n "$( echo $ldrive | grep -E "cciss|ida|nvme|mmcblk" )" ]; then
        # if this is a cciss
        ROOT_PARTITION=$ldrive"p1"
      else
        # else... the rest of the world
        ROOT_PARTITION=$ldrive"1"
      fi
  fi
  # udev takes time to re-add the device file, so wait for it
  while [ ! -b "/dev/$ROOT_PARTITION" ]; do
    sleep 1
  done

  if [ "$initialize_fs" = "yes" ]; then
    # make the root and config file systems.
    make_filesystem "$ROOT_PARTITION"
  fi
}

# ask for user input on the parted and skip setup methods
# $1 is whether or not to run parted
# sets globals INSTALL_DRIVE, ROOT_PARTITION, CONFIG_PARTITION
setup_method_manual() {
  parted=$1

  echo "The VyOS install needs a minimum ${ROOT_MIN}MB root"
  echo "with partiton type 83 (Linux)."
  echo -e "\n\n"

  # if this is parted, let the user create the partitions
  if [ "$parted" == 'parted' ]; then
    while [ -z "$INSTALL_DRIVE" ]; do
      # TODO: right now we only run parted on a single drive
      echo -e "\nI found the following drives on your system:"
      select_drive 'Which drive would you like to run parted on?' \
        'INSTALL_DRIVE'
    done

    # Unmount the install drive if it is mounted
    unmount "$INSTALL_DRIVE"

    # Run parted and let the user configure
    parted /dev/$INSTALL_DRIVE
  fi

  # Ask for the root partition and make sure it's valid
  while [ -z "$ROOT_PARTITION" ]; do
    select_partition 500 "Which partition should I install the root on?"
    # Note that PARTITION is defined in select partition
    ROOT_PARTITION=$PARTITION
    unmount "$ROOT_PARTITION"
    vd=$(grep $ROOT_PARTITION /proc/partitions | awk '{ print $4 }')

    if [ -z "$vd" ]; then
      echo
      echo "$ROOT_PARTITION is an invalid partition. Please try again."
      ROOT_PARTITION=""
    fi
  done

  # See if partition has old data
  mkdir -p /mnt/tmp

  # Is there an old filesystem there?
  if mount -r /dev/"$ROOT_PARTITION" /mnt/tmp 2>&1; then
    save_old_info "$ROOT_PARTITION"
    umount /mnt/tmp
  fi

  # create the filesystem on the part
  make_filesystem "$ROOT_PARTITION"

  # We need to set the INSTALL_DRIVE if it wasn't set when the user ran parted
  # We assume that we will use the boot sector of the same drive that the
  # partition is on.
  # TODO: Allow different drives to function as the boot device
  if [ -z "$INSTALL_DRIVE" ]; then
    if [ ${ROOT_PARTITION:0:2} = "md" ]; then
        INSTALL_DRIVE=$ROOT_PARTITION
    else
        INSTALL_DRIVE=$(echo $ROOT_PARTITION | sed 's/[0-9]//g')
    fi
  fi
}

# Walk the user through the auto setup method
# sets globals INSTALL_DRIVE, ROOT_PARTITION
setup_method_auto () {
  while [ -z "$INSTALL_DRIVE" ]; do
    echo "I found the following drives on your system:"
    select_drive 'Install the image on?' 'INSTALL_DRIVE'

    # check to make sure the drive is large enough to hold the image
    if [ -n "$INSTALL_DRIVE" ]; then
      lsize=$(get_drive_size "$INSTALL_DRIVE")
      total=$ROOT_MIN
      if [ "$total" -gt "$lsize" ]; then
        echo "Unfortunately, VyOS requires a total of at least $total"MB" to properly install."
        echo "$INSTALL_DRIVE is below the minimum required capacity and therefore, cannot be used to"
        echo -e "complete the installation.\n"
        echo "If other drives are present"
        echo -e "Please select another drive...\n"

        INSTALL_DRIVE=''
      fi
    fi
  done

  warn_of_dire_consequences

  echo

  # make sure we aren't working on a mounted part
  unmount "$INSTALL_DRIVE"

  # remove any existing partitions on that drive
  delete_partitions "$INSTALL_DRIVE"

  # Enforce minimum partition size requirement.
  local root_part_size=0
  while [ $ROOT_MIN -gt $root_part_size ]; do
    # Get the size of the drive
    size=$(get_drive_size $INSTALL_DRIVE)
    echo -n "How big of a root partition should I create? ($ROOT_MIN"MB" - $size"MB") [$size]MB: "
    response=$(get_response "$size")
    # TODO: need to have better error checking on this value
    # TODO: This should also probably take into account the size of the EFI partition (256MB)
    root_part_size=$(echo "$response" | sed 's/[^0-9]//g')
    if [ $root_part_size -lt $ROOT_MIN ] \
        || [ $root_part_size -gt $size ]; then
      echo "Root partion must be between $ROOT_MIN"MB" and $size"MB""
      echo
      root_part_size=0
    fi
  done

  echo

  # now take the data and create the partitions
  create_partitions "$INSTALL_DRIVE" "$root_part_size" "yes"
  if ! [ -d /sys/firmware/efi ]; then
      # mark data partition as bootable
      lecho "Marking /dev/$INSTALL_DRIVE partition 1 as bootable"
      output=$(parted -s /dev/$INSTALL_DRIVE set 1 boot on 2>&1)
      lecho "$output"
  fi
  # Must give partition device time to settle
  sleep 5
}

unmount () {
  # grab the list of mounted drives
  # make sure to reverse sort so as to unmount up the tree
  mounted=$(mount | grep "$1" | cut -f3 -d' ' | sort -r)
  if [ -n "$mounted" ]; then
    echo    "I need to unmount: "
    echo    "$mounted"

    response=''
    while [ -z "$response" ]; do
      echo -n "Continue (Yes/No) [No]: "
      response=$(get_response "No" "Yes No Y N")
      if [ "$response" == "no" ] || [ "$response" == "n" ]; then
        echo -e "Ok then.  Need to unmount to continue.\nExiting..."
        exit 1
      fi
    done

    for parts in "$mounted"; do
      lecho "umount $parts"
      output=$(umount $parts)
      status=$?
      if [ "$status" != 0 ]; then
        echo -e "Exiting: error unmounting $parts.\nPlease see $INSTALL_LOG for more details."
        lecho "Exiting: error unmounting $parts.\numount $parts\n$output"
        exit 1
      fi
    done
  fi
}

#####  Main
##
# turn off any mounted swap files
turnoffswap

while true; do

  if ! is_live_cd_boot; then
    # We're running on an installed system, so we don't have to find
    # a partition to install onto
    if is_union_install; then
      # this is a union install
      ROOT_PARTITION_TYPE=union
    else
      # this is an old, non-union install
      ROOT_PARTITION_TYPE=old
    fi
    # flag partition and drive as found but we don't actually need them
    ROOT_PARTITION=dummy
    INSTALL_DRIVE=dummy
    break
  fi

  # some drives don't show up in /proc/partitions so we need to bootstrap them
  echo -n "Probing drives:  "
  progress_indicator start
  probe_drives
  progress_indicator stop
  echo "OK"

  # try old raid
  check_for_old_raid
  if [ -n "$ROOT_PARTITION" ]; then
    # got partition. done.
    break
  fi

  # try new raid
  check_for_new_raid
  if [ -n "$ROOT_PARTITION" ]; then
    # got partition. done.
    break
  fi

  cat <<EOF
The VyOS image will require a minimum ${ROOT_MIN}MB root.
Would you like me to try to partition a drive automatically
or would you rather partition it manually with parted?  If
you have already setup your partitions, you may skip this step

EOF
  PART_METHOD=''
  while [ -z "$PART_METHOD" ]; do
    echo -n "Partition (Auto/Parted/Skip) [Auto]: "
    PART_METHOD=$(get_response "Auto" "Auto Parted Skip A P S")
  done
  echo

  # TODO: Note installs assume an LBA BIOS.  So no boot partition currently.
  # also note that we are not creating a swap partition right now.
  ROOT_PARTITION_TYPE=new
  if [ "$PART_METHOD" == "parted" ] || [ "$PART_METHOD" == "p" ]; then
    setup_method_manual "parted"
  elif [ "$PART_METHOD" == "skip" ] || [ "$PART_METHOD" == "s" ]; then
    setup_method_manual "skip"
  elif [ "$PART_METHOD" == "auto" ] || [ "$PART_METHOD" == "a" ]; then
    setup_method_auto
  fi

  if [ -n "$ROOT_PARTITION" ]; then
    # got partition. done.
    break
  fi

  break
done

if [ -z "$ROOT_PARTITION" ]; then
  echo 'Partition not selected. Exiting...'
  exit 1
fi

echo "$ROOT_PARTITION_TYPE $ROOT_PARTITION $INSTALL_DRIVE" >$OUTFILE
becho 'Done!'
exit 0