#!/bin/sh
#
# Module: install-system
# 
# **** License ****
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
# 
# 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.
#
# This code was originally developed by Vyatta, Inc.
# Portions created by Vyatta are Copyright (C) 2006, 2007 Vyatta, Inc.
# All Rights Reserved.
# 
# Author: Robert Bays
# Date: 2006
# Description:
# 
# **** End License ****
#
# Vyatta system installer script.
#
#
# If you set VYATTA_AUTO_INSTALL I will try to do an automated install for you

INSTALL_LOG="install.log"	# install log file name
ROOT_MIN="450"			# root partition minimum size in MB
CONFIG_MIN="10"			# config partition minium size in MB
VYATTA_CONFIG_DIRECTORY="/opt/vyatta/etc/config/" # the configuration directory

# environmental variables
#$INSTALL_DRIVE			the base install drive e.g. sda
#$ROOT_PARTITION		the install partition e.g. sda1
#$CONFIG_PARTITION		the config partition e.g. sda2
#$ROOT_PARTITION_SIZE		the size of the root partition
#$CONFIG_PARTITION_SIZE		the size of the config partition

# Use "classic" prompts?
CLASSIC=$1

# global holding variable used in the select_partition sub
PARTITION=''

# Process ID of this script for the lame marketing spinner
SPID=$$

# trap signals so we can kill runaway progress indicators
trap 'progress_indicator "stop"; exit 1' 1
trap 'progress_indicator "stop"; exit 1' 2

# turn off any mounted swap partitions
turnoffswap () {
  if [ -f "/proc/swaps" ]; then
    myresponse=$(cat /proc/swaps)
    if [ -n "$myresponse" ]; then
      echo "turning off swaps..." >> $INSTALL_LOG
      swapoff -a
    fi
  fi
}

# Validates a user response.  Returns the response if valid.
# Returns the default is the user just hits enter.  
# Returns nothing if not valid.  Default parameter is $1.
# Options are in $2.  If options are defined return must be a member
# of the enum. 
get_response () {
  ldefault=$(echo "$1" | tr [:upper:] [:lower:])
  loptions=$(echo "$2" | tr [:upper:] [:lower:])

  # get the response from the user
  read myresponse
  myresponse=$(echo "$myresponse" | awk '{ printf (tolower($0)) }')

  # Check to see if the user accepts the default
  if [ -z "$myresponse" ]; then
    echo -n $ldefault
  # if we are passing in options to check, make sure response is a valid option
  elif [ -n "$loptions" ]; then
    for token in $loptions
      do
        if [ "$token" == "$myresponse" ]; then
          echo -n "$myresponse"
          return 0
        fi
      done
    return 1
  else
    echo -n "$myresponse"
  fi

  return 0
}

# Return the size of the drive in MB
get_drive_size () {
  ldrive=$1

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

  # If unable to read disk, it's likely it needs a disklabel
  if [  "$?" != "0" ]; then
    echo "Creating a new disklabel on $ldrive" >> $INSTALL_LOG
    echo "parted /dev/$ldrive mklabel msdos" >> $INSTALL_LOG
    output=$(parted /dev/$ldrive mklabel msdos)

    # Get the drive size from parted
    lsize=$(parted /dev/$ldrive p | grep "^Disk" | awk '{ print $3 }')

    if [ $(echo $lsize | grep error) ]; then
      echo "Unable to read disk label.  Exiting."
      exit 1
    fi
  fi

  # Get the drive size from parted
  lsize=$(parted /dev/$ldrive p | grep "^Disk" | awk '{ print $3 }')
  # Get the reported units (mB, GB, kB)
  lmodifier=$(echo $lsize | sed 's/[0-9\.]//g')
  # remove the modifier
  lsize=$(echo $lsize | sed 's/[a-z,A-Z]//g')
  # Remove any fractions
  lsize=$(echo $lsize | cut -f1 -d'.')
  # Translate our size into mB if not there already
  if [ $lmodifier = "GB" ]; then 
    lsize=$(($lsize * 1000))
  elif [ $lmodifier = "kB" ]; then 
    lsize=$(($lsize / 1000))
  fi

  echo $lsize
}

# Probe hardrives not shown in /proc/partitions by default
probe_drives () {
  # find IDE drives.  Not all drives show up in /proc/partitions for
  # some reason.
  # TODO: not sure we have to do this for anything but compact flash
  # when we test on other platforms, we will find out and modify this section
  # as appropriate
  drive=$(find /proc/ide/ide* -name "hd*" 2>/dev/null | /usr/bin/awk -F"/" '{ print $5 }')

  # now exclude all drives but disks
  for drive in $drive
  do
    media=$(cat /proc/ide/$drive/media)
    if [ "$media" = "disk" ]; then
     output=$(mount | grep $drive)
      if [ -z "$output" ]; then
        output=$(parted /dev/$DRIVE p)
      fi
    fi
  done
}

# Takes an argument to display text before choice
# Sets INSTALL_DRIVE.  Note that select_drive should be wrapped 
# in the verification loop, not the included get_response.
select_drive () {
  # list the drives in /proc/partitions.  Remove partitions and empty lines.
  drives=$(cat /proc/partitions | awk '{ if ($4!="name") { print $4 } }' | egrep -v "[0-9]$" | egrep -v "^$")

  # take the first drive as the default
  INSTALL_DRIVE=$(echo $drives | /usr/bin/awk '{ print $1 }')

  # Add the drive sizes to the display to help the user decide
  display=''
  for drive in $drives
  do
    size=$(get_drive_size $drive)
    display="$display $drive\t$size"MB"\n"
  done

  n=`echo $drives|wc -w`
  if [ $n = 1 -a x$CLASSIC != xCLASSIC ]
  then
    # Only one drive to choose from, no need to ask
    INSTALL_DRIVE=$drives
  else
    # Display the drives and ask the user which one to install to
    echo -e "$display"
    echo
    echo -ne "$1 [$INSTALL_DRIVE]:"

    INSTALL_DRIVE=$(get_response "$INSTALL_DRIVE" "$drives")
    echo
  fi

  # Assume no dma if the disk is smaller than 10G (such as a CF drive)
  size=$(get_drive_size $INSTALL_DRIVE)
  if [ $size -lt 11000 ]
  then
     ISCF="ide=nodma"
  fi
}

# 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

  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')
  
  # Get the partition sizes for display
  # only show partitions that have sizes, i.e. remove loops
  display=''
  for part in $parts
  do
    lsize=$(get_drive_size $part)
    if [ "$lsize" -a $lsize -ge $minsize ]; then
      display="$display $part\t\t$lsize"MB"\n"
    fi
  done

  progress_indicator "stop" 
  echo "OK"

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

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

      lpartition=$(get_response "$lpartition" "$parts")
      echo
    done
  else
    echo "No suitable partition sizes found.  Exiting..."  | tee $INSTALL_LOG
    exit 1
  fi
  PARTITION=$lpartition
}

# 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
  partitions=$(cat /proc/partitions | grep $ldrive[0-9] | awk '{ print $4 }' | sed 's/[a-z]//g')

  # now for each part, blow it away
  for part in $partitions
  do
    # Look to see if this is a config partition
    mkdir -p /mnt/tmp
    output=$(mount /dev/$ldrive$part /mnt/tmp 2>&1)
    if [ $? != 0 ]; then
      echo -e "Cannot mount /dev/$ldrive$part"."\nPlease see $INSTALL_LOG for more details.\nExiting.."
      echo -e "Cannot mount /dev/$ldrive$part"."\nmount /dev/$ldrive$part /mnt/tmp\nExiting..." >> $INSTALL_LOG
      echo "$output" >> $INSTALL_LOG
      exit 1
    fi

    # Look to see if the old config marker file is there.
    if [ -f /mnt/tmp/.vyatta_config ]; then
      response=''
      while [ -z "$response" ]
      do
        echo "/dev/$ldrive$part is a config partition!"
        echo -ne "Would you like me to save the data on it before I delete it? (Yes/No) [Yes]: "
        response=$(get_response "Yes" "Yes No Y N")
        if [ "$response" == "yes" ] || [ "$response" == "y" ]; then
          mkdir -p /mnt/config
          output=$(cp -pR /mnt/tmp/* /mnt/config)
          if [ -n "$output" ]; then
            echo -e "Warning: error in copying the old config partition.\nSee $INSTALL_LOG for more details."
            echo -e "Warning: error in copying the old config partition.\ncp -pR /mnt/tmp/* /mnt/config\n$output\n" >> $INSTALL_LOG
          fi
        fi
      done
    fi  
    umount /mnt/tmp

    echo "Removing partition $part on /dev/$ldrive" >> $INSTALL_LOG
    output=$(parted /dev/$ldrive rm $part)
    status=$?
    if [ "$status" != 0 ]; then
      echo -e "Warning: cannot delete partition $part on $ldrive.\nPlease see $INSTALL_LOG for more details."
      echo -e "Warning: cannot delete partition $part on $ldrive.\nparted /dev/$ldrive rm $part\n$output" >> $INSTALL_LOG
    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:  "
  echo "Creating filesystem on /dev/$ldrive..." >> $INSTALL_LOG

  progress_indicator "start"
  output=$(mke2fs -j /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..."
    echo -e "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
# $3 config partition size e.g. 100
# This will set the globals ROOT_PARTITION and CONFIG_PARTITION
create_partitions() {
  ldrive=$1
  root_part_size=$2
  config_part_size=$3

  total=$(($root_part_size + $config_part_size))
  head=$(($root_part_size))
  tail=$(($root_part_size + $config_part_size))

  # Make sure there is enough space on drive
  size=$(get_drive_size "$ldrive") 
  if [ "$total" -gt "$size" ]; then
    echo "Error: $ldrive is only $size"MB" large."
    exit 1
  fi

  echo "Creating root partition on /dev/$ldrive" >> $INSTALL_LOG
  
  # make the root partition
  output=$(parted /dev/$ldrive mkpart primary 0 $root_part_size)
  status=$?
  if [ "$status" != 0 ]; then
    echo -e "Error creating primary partition on $ldrive.\nPlease see $INSTALL_LOG for more details.\nExiting..."
    echo -e "Error creating primary partition on $ldrive.\nparted /dev/$ldrive mkpart primary 0 $root_part_size\n$output" >> $INSTALL_LOG
    exit 1
  fi

  ROOT_PARTITION=$ldrive"1"
  # udev takes time to re-add the device file, so wait for it
  while [ ! -b "/dev/$ROOT_PARTITION" ] 
  do 
    sleep 1
  done

  echo "Creating configuration partition on /dev/$ldrive" >> $INSTALL_LOG
  echo "parted /dev/$ldrive mkpart primary $root_part_size $config_part_size" >> $INSTALL_LOG

  output=$(parted /dev/$ldrive mkpart primary $head $tail)
  status=$?
  if [ "$status" != 0 ]; then
    echo -e "Error creating primary partition on $ldrive.\nPlease see $INSTALL_LOG for more details.\nExiting..."
    echo -e "Error creating primary partition on $ldrive.\nparted /dev/$ldrive mkpart primary $root_part_size $config_part_size\n$output"
    exit 1
  fi

  CONFIG_PARTITION=$ldrive"2"
  # udev takes time to re-add the device file, so wait for it
  while [ ! -b "/dev/$CONFIG_PARTITION" ]
  do
    sleep 1
  done

  # make the root and config ext3 file systems.
  make_filesystem "$ROOT_PARTITION"
  make_filesystem "$CONFIG_PARTITION"
}

# Install the root filesystem
# $1 is the partition to install on
install_root_filesystem () {
  ldrive=$1

  dpkg="/usr/bin/dpkg --force-all --root=/mnt/rootfs"
  echo "Mounting /dev/$ldrive..." | tee $INSTALL_LOG

  # make the mount point
  output=$(/bin/mkdir -p /mnt/rootfs)

  # mount the partition
  output=$(mount -t ext3 /dev/$ldrive /mnt/rootfs)
  status=$?
  
  if [ "$status" != 0 ]; then
    echo -e "Error trying to mount the new root partition.\nPlease see $INSTALL_LOG for details.\nExiting..."
    echo -e "Error trying to mount the new root partition.\nmount -t ext3 /dev/$ldrive /mnt/rootfs\n$output" >> $INSTALL_LOG
    exit 1
  fi
  
  echo -n "Copying system image files to /dev/$ROOT_PARTITION: "
  progress_indicator "start"
  # Mount the squashfs for copying
  output=$(mkdir -p /mnt/squashfs)
  output=$(mount /live_media/casper/filesystem.squashfs /mnt/squashfs -t squashfs -o loop)
  status=$?
  
  if [ "$status" != 0 ]; then
    echo -e "Error trying to mount the squashfs.\nPlease see install log for more details.\nExiting..."
    echo -e "Error trying to mount the squashfs.\nmount /live_media/casper/filesystem.squashfs /mnt/squashfs -t squashfs -o loop\n$output" >> $INSTALL_LOG
    exit 1
  fi

  output=$(cp -pR /mnt/squashfs/* /mnt/rootfs/)
  status=$?
  
  if [ "$status" != 0 ]; then
    echo -e "Error trying to copy the rootfs.\nPlease see install log for more details.\nExiting.
.."
    echo -e "Error trying to copy the rootfs.\ncp -pR /mnt/squashfs/* /mnt/rootfs/\n$output" >> $INSTALL_LOG
    exit 1
  fi

  # unmount the squashfs.  No big deal if it fails.
  output=$(umount /mnt/squashfs)

  # make the dir for the boot files and copy em' over
  mkdir -p /mnt/rootfs/boot
  output=$(cp -pR /boot/* /mnt/rootfs/boot/)
  status=$?
  output+=$(cp /live_media/boot/initrd.img /mnt/rootfs/boot/)
  status=$status || $?
  if [ "$status" != 0 ]; then
    echo -e "Error trying to copy the bootfiles.\nPlease see install log for more details.\nExiting.
.."
    echo -e "Error trying to copy the bootfiles.\ncp /live_media/boot/initrd.img /mnt/rootfs/boot/\n$output" >> $INSTALL_LOG
    exit 1
  fi

  # create the fstab
  echo -e "/dev/$CONFIG_PARTITION\t/opt/vyatta/etc/config\text3\tdefaults\t1 2" >> /mnt/rootfs/etc/fstab
  echo -e "/dev/$ROOT_PARTITION\t/\text3\tdefaults\t0 1" >> /mnt/rootfs/etc/fstab
  progress_indicator "stop"
  echo "OK"
}

# copy the configuration to the config partition
# $1 is the config partition device
copy_config () {
  config_partition=$1
  lerror=''

  # Copy the config.boot.default to tmp
  output=$(cp $VYATTA_CONFIG_DIRECTORY/config.boot.default /tmp/config.boot.default.install)

  # create the config directory on the union file system
  output=$(mount -t ext3 /dev/$config_partition /mnt/rootfs/$VYATTA_CONFIG_DIRECTORY)
  status=$?

  if [ "$status" != 0 ]; then
    echo -e "Error trying to mount the new config partition.\nPlease see $INSTALL_LOG for details.\nExiting..."
    echo -e "Error trying to mount the new config partition.\nmount -t ext3 /dev/$config_partition /mnt/rootfs/$VYATTA_CONFIG_DIRECTORY\n$output" >> $INSTALL_LOG
    exit 1
  fi

  # create the proper perms on the new config partition
  chgrp xorp /mnt/rootfs/$VYATTA_CONFIG_DIRECTORY
  chmod 775 /mnt/rootfs/$VYATTA_CONFIG_DIRECTORY

  # Copy /tmp/config.boot.default to new config partition
  cp /tmp/config.boot.default.install /mnt/rootfs/$VYATTA_CONFIG_DIRECTORY/config.boot.default
  rm -f /tmp/config.boot.default.install

  # create our config partition marker
  touch /mnt/rootfs/$VYATTA_CONFIG_DIRECTORY/.vyatta_config

  if [ -d /mnt/config ]; then
    echo "Copying old configurations to config partition."
    cp -pR /mnt/config/* /mnt/rootfs/$VYATTA_CONFIG_DIRECTORY
  else
    # Find the config files and give the user the option to copy config files
    # TODO: this needs cleaned up
    [ -f $VYATTA_CONFIG_DIRECTORY/config.boot ] &&  config="$VYATTA_CONFIG_DIRECTORY/config.boot"
    fd_config=/media/floppy/config/config.boot
    if [ ! -f $fd_config ]; then
	configfile="$config"
    else
      if [ "$config" ]; then
        config="$config\n$fd_config"
      else
        config="$fd_config"
	configfile="$config"
      fi
    fi

    if [ "$config" -a ! $configfile ]; then
      echo "I found the following configuration files:"
      echo -e "$config"
      default=$(echo -e $config| head -1)

      resp=''
      while [ -z "$configfile" ]
      do
        echo -n "Which one should I copy to $INSTALL_DRIVE? [$default]: "
        configfile=$(get_response "$default" "$config")
      done
      
      echo
      if ! cp $configfile /mnt/rootfs/$VYATTA_CONFIG_DIR; then
        echo "Error copying file $configfile to config directory.  Exiting..." >> $INSTALL_LOG
        exit 1
      fi
    fi
  fi

  # set the permissions on the new config file
  if [ -f "/mnt/rootfs/$VYATTA_CONFIG_DIR/config.boot" ]; then
    chgrp xorp /mnt/rootfs/$VYATTA_CONFIG_DIR/config.boot
    chmod 775 /mnt/rootfs/$VYATTA_CONFIG_DIR/config.boot
  fi
}

# setup grub on thje boot sector of a user queried drive
install_grub () {
  orig_install_drive="$INSTALL_DRIVE"
  # we now use INSTALL_DRIVE to reference the grub boot drive.
  #  that way I can re-use select_drive.  I'm lazy that way.
  INSTALL_DRIVE=''

  mkdir -p /mnt/rootfs/boot/grub
  # Let the user choose the boot sector
  
  while [ -z "$INSTALL_DRIVE" ]
  do
    select_drive "I need to install the GRUB bootloader.\nI found the following drives on your system:\nWhich drive should GRUB modify the boot partition on?"
  done

  echo -n "Setting up grub:  "
  echo "Setting up grub..." >> $INSTALL_LOG

  # Install grub in the boot sector of the primary drive
  progress_indicator "start"
  grub-install --no-floppy --root-directory=/mnt/rootfs /dev/$INSTALL_DRIVE >>$INSTALL_LOG 2>&1
  progress_indicator "stop"

  # TODO: This needs to be changed to map to the correct drive
  part=$(echo $ROOT_PARTITION | sed 's/[^0-9]//g')
  part=$(($part - 1))
  if [ $(cat /sys/block/$orig_install_drive/removable) == 0 ]; then
    root=$(grep $orig_install_drive /mnt/rootfs/boot/grub/device.map | /usr/bin/awk -F')' '{ print $1 }')
    root="$root,$part)"
  else 
    echo "This looks like a removable device.  Setting root grub device to (0,0)." | tee $INSTALL_LOG
    root="(hd0,$part)"
  fi
  
  # create the menu.lst file for grub
  grub_file="/mnt/rootfs/boot/grub/menu.lst"
  echo -e "default=0\ntimeout=5\n#splashimage=(hd0,0)/grub/splash.xpm.gz\nhiddenmenu" > $grub_file
  # set serial console options
  echo -e "serial --unit=0 --speed=9600\nterminal --timeout=5 console serial\n\n" >> $grub_file
  # set primary boot option
  echo -e "title Vyatta OFR\n\troot $root" >> $grub_file
  echo -en "\tkernel /boot/vmlinuz root=/dev/$ROOT_PARTITION $ISCF console=ttyS0,9600 console=tty0\n" >> $grub_file
  echo -e "\tinitrd /boot/initrd.img" >> $grub_file

  echo "OK"
}

# 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 Vyatta install needs a minimum $ROOT_MIN"MB" root"
  echo "and a minimum $CONFIG_MIN"MB" config partition. These"
  echo "partitions need to be set to partiton type 83 (Linux)."
  echo -e "\n\n"

  # if this is parted, let the user create the partitions
  if [ "$method" == "parted" ] || [ "$method" == "p" ]; then
    while [ -z "$INSTALL_DRIVE" ]
    do
      # TODO: right now we only run parted on a single drive
      select_drive "\nI found the following drives on your system:\nWhich drive would you like to run parted on?"
    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

  # Ask for the config partition and make sure it's valid
  # TODO: need to do better error checking here to insure
  # the user doesn't select the same part for both drives.
  while [ -z "$CONFIG_PARTITION" ]
  do
    while true
    do
        select_partition 5 "Which partition should I use for configurations?"
        # Note that PARTITION is defined in select partition
	if [ "$PARTITION" != "$ROOT_PARTITION" ]
        then
	    break
	fi
	echo "The config partition cannot be the same as the root partition ($ROOT_PARTITION)!"
	echo
    done

    # Note that PARTITION is defined in select partition
    CONFIG_PARTITION=$PARTITION
    unmount "$CONFIG_PARTITION"
    vd=$(grep $CONFIG_PARTITION /proc/partitions | awk '{ print $4 }')

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

  # create the ext3 fs on the part
  make_filesystem "$ROOT_PARTITION"

  # Create the ext3 fs on the part
  make_filesystem $CONFIG_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
    INSTALL_DRIVE=$(echo $ROOT_PARTITION | sed 's/[0-9]//g')
  fi
}

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

    # 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 + $CONFIG_MIN))
      if [ "$total" -gt "$lsize" ]; then
        echo "Unfortunately, the OFR 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

  # Give the user a requisite warning that we are about to nuke their drive
  response=''
  while [ -z $response ]
  do
    echo -n "This will destroy all data on /dev/$INSTALL_DRIVE. 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

  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 partion size requirement.
  ROOT_PARTITION_SIZE=0
  while [ $ROOT_MIN -gt $ROOT_PARTITION_SIZE ]; do
    # Get the size of the drive
    size=$(get_drive_size $INSTALL_DRIVE)
    # Subtract our minimum config part so the user doesn't fill the entire drive
    size=$(($size - $CONFIG_MIN))
    if [ $ROOT_MIN = $size ]
    then
	response=$size
    else
        echo -ne "\nHow big of a root partition should I create? ($ROOT_MIN"MB" - $size"MB") [$size]MB: "
        response=$(get_response "$size")
    fi

    # TODO: need to have better error checking on this value
    ROOT_PARTITION_SIZE=$(echo "$response" | sed 's/[^0-9]//g')
    if [ $ROOT_PARTITION_SIZE -lt $ROOT_MIN ] || [ $ROOT_PARTITION_SIZE -gt $size ]; then
      echo "Root partion must be between $ROOT_MIN"MB" and $size"MB""
      echo
      ROOT_PARTITION_SIZE=0
    fi
  done

  # Enforce minimum partion size requirement.
  CONFIG_PARTITION_SIZE=0
  while [ $CONFIG_MIN -gt $CONFIG_PARTITION_SIZE ]; do
    # Get the size of the drive
    size=$(get_drive_size $INSTALL_DRIVE)
    # Subtract our minimum config part so the user doesn't fill the entire drive
    size=$(($size - $ROOT_PARTITION_SIZE))
    if [ $CONFIG_MIN = $size ]
    then
	response=$size
    else
	echo -n "How big of a config partition should I create? ($CONFIG_MIN"MB" - $size"MB") [$CONFIG_MIN]MB: "
	response=$(get_response "$CONFIG_MIN")
    fi
    CONFIG_PARTITION_SIZE=$(echo "$response" | sed 's/[^0-9]//g')
    if [ $CONFIG_PARTITION_SIZE -lt $CONFIG_MIN ] || [ $CONFIG_PARTITION_SIZE -gt $size ]; then
      echo "Config partion must be between $CONFIG_MIN"MB" and $size"MB""
      echo
      CONFIG_PARTITION_SIZE=0
    fi
  done

  echo

  # now take the data and create the partitions
  create_partitions "$INSTALL_DRIVE" "$ROOT_PARTITION_SIZE" "$CONFIG_PARTITION_SIZE"
}

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
      echo "umount $parts" >> $INSTALL_LOG
      output=$(umount $parts)
      status=$?
      if [ "$status" != 0 ]; then
        echo -e "Exiting: error unmounting $parts.\nPlease see $INSTALL_LOG for more details."
        echo -e "Exiting: error unmounting $parts.\numount $parts\n$output" >> $INSTALL_LOG
        exit 1
      fi
    done
  fi
}

progress_indicator () {
  case "$1" in
    "start") /usr/bin/progress-indicator $SPID &
             ;;
     "stop") rm -f /tmp/pi.$SPID
             sleep 1
	     echo -n -e "\b"
             ;;
  esac
}

#####  Main
##
# clean up existing log files
if [ -f $INSTALL_LOG.old ]; then 
  rm -f $INSTALL_LOG.old
fi
if [ -f $INSTALL_LOG ]; then
  mv $INSTALL_LOG $INSTALL_LOG.old
  rm -f $INSTALL_LOG
fi

# turn off any mounted swap files
turnoffswap

# Print welcome and instructions.
echo "Welcome to the Vyatta install program.  This script will walk you through"
echo "the process of installing the Vyatta image to a local hard drive."
echo

if [ x$CLASSIC = xCLASSIC ]
then
	response=''
	while [ -z $response ]
	do
	  echo -n "Would you like to continue? (Yes/No) [Yes]: "
	  response=$(get_response "Yes" "Yes No Y N")
	  if [ "$response" == "no" ] || [ "$response" == "n" ]; then
	    echo "Ok then.  Exiting..."
	    exit 1
	  fi
	done
fi

echo "The image will require a minimum $ROOT_MIN"MB" root partition and"
echo "a minimum $CONFIG_MIN"MB" configuration partition."
echo
echo "Disk partitioning can be done automatically or manually (with parted)."
echo "If you have already set up your partitions, you may skip the Partition step."
echo

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

method=''
while [ -z $method ]
do
  echo -n "Partition (Auto/parted/skip) [Auto]: "
  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.
if [ "$method" == "parted" ] || [ "$method" == "p" ]; then
  setup_method_manual "parted"
elif [ "$method" == "skip" ] || [ "$method" == "s" ]; then
  setup_method_manual "skip" 
elif [ "$method" == "auto" ] || [ "$method" == "a" ]; then
  setup_method_auto
elif [ "$method" == "vyatta" ]; then
  echo "Automated install..."
  echo "unmounting $INSTALL_DRIVE"
  unmount "$INSTALL_DRIVE"
  echo "deleting partitions on $INSTALL_DRIVE"
  delete_partitions "$INSTALL_DRIVE"
  echo "creating config partition"
  create_partitions "$INSTALL_DRIVE" "$ROOT_PARTITION_SIZE" "$CONFIG_PARTITION_SIZE"
fi

# Install the root filesystem
install_root_filesystem "$ROOT_PARTITION"

# Copy the config files
copy_config "$CONFIG_PARTITION"

# Install grub
install_grub

cp $INSTALL_LOG /mnt/rootfs/install.log

umount /mnt/rootfs/opt/vyatta/etc/config
umount /mnt/rootfs

echo "Done!" | tee $INSTALL_LOG
echo "Don't forget to remove the live CD after system shutdown"

echo -ne "Reboot the system now? (Yes/No) [Yes]: "
response=$(get_response "Yes" "Yes No Y N")
if [ $response = yes -o $response = y ]
then
  reboot
fi

exit 0