#!/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
}

delete_old_raid () {
  ldrive=$1

  # 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 --run /dev/md?* >/dev/null 2>&1
  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
    return
  fi

  for drive in $raid_drives; do
    members=`ls /sys/block/$drive/slaves`
    for member in $members; do
      cdrive=`lsblk -dn -o pkname /dev/$member`
      if [ "$ldrive" == "$cdrive" ]; then
        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}:"
          if [ -d /sys/firmware/efi ]; then
            sgdisk --part-type /dev/$drive $part 0x83
          else
            sfdisk --part-type /dev/$drive $part 0x83
          fi
            echo "Clearing RAID superblock from RAID group $raid_drive member /dev/${member}."
            mdadm --zero-superblock /dev/$member
        done
      fi
    done
  done
}

check_for_new_raid () {
  # Identify physical drives
  driveName=()
  driveSize=()
  driveNameSize=()
  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"/}
      else
          driveName+=("$instdrv")
          driveSize+=($(get_drive_size $instdrv))
          driveNameSize+=("$instdrv $(get_drive_size $instdrv) MB")
      fi
  done

  numdrives=`echo $drives | wc -w`

  # Need at least two drives for RAID-1. If there are more than two
  # drives found we have the option to check for two drives of the same size
  # or configure manually.

  if [ $numdrives -lt 2 ]; then
    return
  elif [ $numdrives -gt 2 ]; then
    echo "Found ${numdrives} drives."
    echo -n "Would you like to configure RAID-1 mirroring? (Yes/No) [Yes]:"
    response=$(get_response "Yes" "Yes No Y N")
    if [ "$response" == "no" ] || [ "$response" == "n" ]; then
      echo "Ok.  Not configuring RAID-1."
      return
    else
      echo -n "Would you like find and use 2 drives of the same size? (Yes/No) [Yes]:"
      response=$(get_response "Yes" "Yes No Y N")
      if [ "$response" == "yes" ] || [ "$response" == "y" ]; then
        raiddrives=$(get_size_raid)
        if [ -n "$raiddrives" ]; then
          drives=$raiddrives
        else
          echo "No drives with the same size found."
          echo -n "Would you like to manually select the drives? (Yes/No) [Yes]:"
          response=$(get_response "Yes" "Yes No Y N")
          if [ "$response" == "yes" ] || [ "$response" == "y" ]; then
            raiddrives=$(get_manual_raid)
            if [ -n "$raiddrives" ]; then
              drives=$raiddrives
	    fi
          else
            echo "Ok.  Not configuring RAID-1."
            return
          fi
        fi
      else
        echo -n "Would you like to manually select the drives? (Yes/No) [Yes]:"
        response=$(get_response "Yes" "Yes No Y N")
        if [ "$response" == "yes" ] || [ "$response" == "y" ]; then
          raiddrives=$(get_manual_raid)
          if [ -n "$raiddrives" ]; then
            drives=$raiddrives
	  fi
        else
          echo "Ok.  Not configuring RAID-1."
          return
        fi
      fi
    fi
  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 "No" "Yes No Y N")
  if [ "$response" == "no" ] || [ "$response" == "n" ]; then
    echo "Ok.  Not configuring RAID-1."
    return
  fi

  config_saved='no'
  for drive in $drives; do
    echo "Deleting old partitions on drive $drive"
    # remove any existing partitions on that drive
    delete_partitions "$drive" "$config_saved"
    config_saved='yes'
  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 --part-type /dev/$drive $data_dev 0xfd >/dev/null 2>&1
        # 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/'
    cp /usr/share/initramfs-tools/scripts/local-block/mdadm /etc/initramfs-tools/scripts/local-top/
	sed -i 's/\$((COUNT + 1))/20/g' /etc/initramfs-tools/scripts/local-top/mdadm
    if [ -e /usr/sbin/update-initramfs ]; then
      echo "Updating initramfs to include the raid config:"
      /usr/sbin/update-initramfs -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
# $2 is to save the config from only the first drive
delete_partitions () {
  ldrive=$1
  config_saved=$2

  delete_old_raid $ldrive

  # 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
            if [ "$config_saved" == "no" ]; then
	        save_old_info $ldrive$lpart
            fi
	    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
  config_saved='no'
  delete_partitions "$INSTALL_DRIVE" "$config_saved"

  # 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 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