#!/bin/bash

# this script installs a new release image into a running system.
# the new ISO image must be already mounted by caller.
# the script sets up a new union mount for the new release. a reboot
# is then required to boot into the newly installed release.

# the current install type: "union" or "old"
CUR_INSTALL=$1

source /opt/vyatta/sbin/install-functions

failure_exit () {
  echo "$*"
  exit 1
}

# Determine the default menuentry index for the grub.cfg file
get_grub_index () {
    cur_index=`grep "^set default=" $BOOT_DIR/grub/grub.cfg | \
	awk -F= '{print $2}'`

    (( cur_index++ ))

    cur_line=`grep "^menuentry" $BOOT_DIR/grub/grub.cfg | \
	tail -n +${cur_index} - | head -n 1 -`

    is_kvm=`echo $cur_line | grep KVM`
    is_serial=`echo $cur_line | grep Serial`

    # index 0 is KVM, 1 is Serial
    if [ -n "$is_serial" ]; then
	echo 1
    else
	echo 0
    fi
}

# Check if installing on AWS EC2 AMI
is_amazon_ec2_ami () {
  /opt/vyatta/sbin/ec2-check.pl
  if [ $? != 0 ]; then
      return 1
  fi

  ami_id_url=http://169.254.169.254/latest/meta-data/ami-id

  ami_id=$(/usr/bin/curl --silent "$ami_id_url")
  if [ -n "$ami_id" ]; then
    echo "Installing on VyOS AMI"

    # Create init script links for /etc/init.d/ec2-vyos-init
    chroot $INST_ROOT update-rc.d ec2-vyos-init defaults &>/dev/null

    # Dijkstra, forgive us!
    return 0
  else
    return 1
  fi
}

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

# On image-installed systems, the image name can be found as the
# directory under "/boot" on the path to the running kernel on the
# boot line. On non-image-installed systems, this yelds the
# name of the kernel image file.
CURVER=`awk '{print $1}' /proc/cmdline`
CURVER=${CURVER#BOOT_IMAGE=/boot/}
CURVER=${CURVER%/vmlinuz*}

if [ -z "$CURVER" ]; then
  failure_exit 'Cannot find current version.'
fi

# get new version string. this is from the squashfs image.
NEWVER=`cat ${CD_SQUASH_ROOT}/opt/vyatta/etc/version | grep "Version:" | awk '{print $2,$3}' | sed 's/[[:space:]]*$//' | sed 's/ /-/g'`
NEWNAME=$NEWVER

echo -n "What would you like to name this image? [$NEWNAME]: "
if [ "$VYATTA_PROCESS_CLIENT" == "gui2_rest" ]; then
    response=$NEWNAME
else
    read response
fi
if [ -n "$response" ]; then
    badchars=`echo $response | sed -e 's/[a-zA-Z0-9\.\_+-]//g'`
    if [ -n "$badchars" ]; then
	echo "Image name must be composed of a-z, A-Z, 0-9, or one of ._+-"
	exit 1
    fi
    NEWNAME=$response
fi

if [ -z "$NEWNAME" ]; then
  failure_exit 'Invalid image name.'
fi

# Validate image name
if [ "$NEWNAME" = "grub" -o "${NEWNAME:0:7}" = "vmlinuz" -o \
    "${NEWNAME:0:6}" = "initrd" -o "${NEWNAME:0:10}" = "System.map" -o \
    "$NEWNAME" = "Old-non-image-installation" ]; then
    echo "Can't use $NEWNAME.  It is a reserved image name."
    exit 1;
fi

echo "OK.  This image will be named: $NEWNAME"

# this is the default if current install is union

if [ "$CUR_INSTALL" == 'old' ]; then
  BOOT_DIR=/boot
elif [ "$CUR_INSTALL" == 'union' ]; then
  BOOT_DIR=`/opt/vyatta/sbin/vyos-persistpath`/boot
else
  echo 'Invalid current install type. Exiting...'
  exit 1
fi

if [ -d $BOOT_DIR/$NEWNAME ]; then
    if [ "$CURVER" = "$NEWNAME" ]; then
           echo "$NEWNAME is the image you are currently running.  Can't"
	   echo "Re-install over the running image."
	   exit 1
    fi

    echo "An image named $NEWNAME is already installed on this system."
    echo "Proceeding with this installation will delete this copy of"
    echo "$NEWNAME and replace it with a new copy."
    echo -n "Do you want to replace it? (Yes/No) [No]: "
    resp=$(get_response "No" "Yes No Y N")
    if [ "$resp" != 'yes' ] && [ "$resp" != 'y' ]; then
	echo "OK.  Will not replace $NEWNAME"
	echo "Exiting..."
	exit 1
    fi
fi

#
# Check to make sure we have enough space to install image...
#
space_avail=`df -k / | tail -1 | awk '{ print $4 }'`
space_needed=`du -s ${CD_ROOT}/live | awk  '{ print $1 }'`

if [ $space_avail -le $space_needed ]; then
    echo "We do not have enough disk space to install this image!"
    echo "We need $space_needed KB, but we only have $space_avail KB."
    echo "Exiting..."
    exit 1
fi

# start the install
echo "Installing \"$NEWNAME\" image."

# create the new release directories
REL_ROOT=$BOOT_DIR/$NEWNAME
RW_DIR="$REL_ROOT/rw"
if ! mkdir -p "$RW_DIR"; then
  failure_exit 'Cannot create directory for new release.'
fi

WORK_DIR="$REL_ROOT/work"
mkdir -p "$WORK_DIR"

# copy the squashfs image and boot files
echo "Copying new release files..."
squash_img=${CD_ROOT}/live/filesystem.squashfs
boot_dir=${CD_SQUASH_ROOT}/boot
boot_files=$(find $boot_dir -maxdepth 1 -type f -o -type l 2>/dev/null)
if [ ! -f "$squash_img" ] || [ -z "$boot_files" ]; then
  becho 'Cannot find the files. Exiting...'
  exit 1
fi
target_squash=$REL_ROOT/$NEWVER.squashfs
cp -p $squash_img $target_squash >&/dev/null
cp --no-dereference --preserve=all $boot_files $REL_ROOT/ >&/dev/null

# mount copied squashfs
if ! try_mount "-o loop,ro $target_squash $READ_ROOT"; then
  rm -rf $REL_ROOT
  failure_exit 'Failed to mount new squashfs image.'
fi

# set up root for postinst
margs=$(gen_mopts "overlay" $RW_DIR $READ_ROOT $WORK_DIR $INST_ROOT)
if ! try_mount "$margs"; then
  rm -rf $REL_ROOT
  failure_exit 'Failed to set up root directory for postinst.'
fi

#
# Check to make sure we have enough space to copy the config and data dirs...
#
space_avail=`df -k / | tail -1 | awk '{ print $4 }'`
if [ -e $${VYATTA_NEW_CFG_DIR}/data ]; then
    space_needed_data=`du -s ${VYATTA_NEW_CFG_DIR}/data | awk  '{ print $1 }'`
else
    space_needed_data=0
fi
space_needed_configdata=`du -s ${VYATTA_NEW_CFG_DIR} | awk  '{ print $1 }'`
space_needed_config=$(($space_needed_configdata - $space_needed_data))

DEMO_MIGRATE_SCRIPT=/opt/vyatta/sbin/demo-to-vse.pl

# save current config dir if needed
if [ $space_avail -gt $space_needed_configdata ]; then
  resp=''
  while [ -z "$resp" ]; do
    echo 'Would you like to save the current configuration '
    echo -n 'directory and config file? (Yes/No) [Yes]: '
    resp=$(get_response "Yes" "Yes No Y N")
    if [ "$resp" == 'yes' ] || [ "$resp" == 'y' ]; then
      echo 'Copying current configuration...'
      ndir=${INST_ROOT}/${VYATTA_NEW_CFG_DIR}
      mkdir -p $ndir
      find $VYATTA_NEW_CFG_DIR -maxdepth 1 -mindepth 1 \
        -exec cp '-a' '{}' "$ndir/" ';'
      chgrp -R vyattacfg $ndir
      chmod -R 775 $ndir
      if [ -x $DEMO_MIGRATE_SCRIPT ]; then
	  $DEMO_MIGRATE_SCRIPT $ndir/config.boot
      fi
    fi
  done
else
    echo 'There is not enough space to save the current configuration directory.'
    echo -n 'Would you like to continue without saving it? (Yes/No) [No]: '
    resp=$(get_response "No" "Yes No Y N")
    if [ "$resp" == 'yes' ] || [ "$resp" == 'y' ]; then
	echo "OK. Proceeding without saving current config directory."
    else
	rm -rf $REL_ROOT
	failure_exit "OK.  Exiting."
    fi
fi

if [ -d /etc/ssh ]; then
  resp=''
  while [ -z "$resp" ]; do
    echo 'Would you like to save the SSH host keys from your '
    echo -n 'current configuration? (Yes/No) [Yes]: '
    resp=$(get_response "Yes" "Yes No Y N")
    if [ "$resp" == 'yes' ] || [ "$resp" == 'y' ]; then
      echo 'Copying SSH keys...'
      ndir=${INST_ROOT}/etc/ssh
      mkdir -p $ndir
      cp -p /etc/ssh/ssh_host* $ndir
    fi
  done
fi

# postinst hook
PI_SCRIPT=${INST_ROOT}${vyatta_sysconfdir}/install-image/postinst
if [ -e "$PI_SCRIPT" ]; then
  echo "Running post-install script..."
  $PI_SCRIPT $INST_ROOT
fi

# set up grub entry (if provided)
DEF_GRUB=${INST_ROOT}${vyatta_sysconfdir}/grub/default-union-grub-entry
if [ -e "$DEF_GRUB" ]; then
  echo "Setting up grub configuration..."
  new_index=$(get_grub_index)

  def_grub_vers=/tmp/def_grub.$$
  cp $DEF_GRUB $def_grub_vers
  sed -i "s/menuentry \"VyOS.*(/menuentry \"VyOS $NEWNAME (/" $def_grub_vers
  sed -i "s/menuentry \"Lost password change.*(/menuentry \"Lost password change $NEWNAME (/" $def_grub_vers
  sed -i "sX/boot/[A-Za-z0-9\.\-]*X/boot/${NEWNAME}Xg" $def_grub_vers

  old_grub_cfg=$BOOT_DIR/grub/grub.cfg
  new_grub_cfg=/tmp/grub.cfg.$$
  sed -n '/^menuentry/q;p' $old_grub_cfg >$new_grub_cfg
  cat $def_grub_vers >> $new_grub_cfg
  sed -n '/^menuentry/,${p}' $old_grub_cfg >>$new_grub_cfg
  sed -i "s/^set default=[0-9]\+$/set default=$new_index/" $new_grub_cfg
  mv $new_grub_cfg $old_grub_cfg

  # Update the default image symlink used by Xen
  if [ -L $BOOT_DIR/%%default_image ]; then
      mv $BOOT_DIR/%%default_image $BOOT_DIR/%%default_image.orig
      ln -s $NEWNAME $BOOT_DIR/%%default_image
  fi

  # Modify grub.cfg for AWS EC2 AMI
  if is_amazon_ec2_ami; then
    sed -i "/menuentry \"VyOS $NEWNAME (Serial/{N;N;N;N;d;}" $BOOT_DIR/grub/grub.cfg
    sed -i "/menuentry \"Lost password change $NEWNAME/{N;N;N;N;d;}" $BOOT_DIR/grub/grub.cfg
    sed -i "s/VyOS $NEWNAME (KVM console)/VyOS AMI (HVM) $NEWNAME/" $BOOT_DIR/grub/grub.cfg
    sed -i "s/$NEWNAME console=ttyS0.*/$NEWNAME console=ttyS0/" $BOOT_DIR/grub/grub.cfg
  fi
fi

logger -p local3.warning -t "SystemImage" "System Image $NEWNAME has been added and made the default boot image"

echo 'Done.'

# done
exit 0