#!/bin/bash -e
#
# Authors:
#   Tommy Chang (chang177@purdue.edu)
#
# Version 2017-12-16
#
# Assumptions:
#
#   We assume that your new server node has no OS and has a pair of
#   matching harddrives.
#
# Outcome:
#
#   After running this script on a USB-booted (or PXE-booted) Linux,
#   your server computer will be ready to boot to the Ubuntu OS that
#   you provided.  Additionally, the script also installs Free DOS
#   (good for updating BIOS and other firmwares), as well as Gparted
#   Linux.  Gparted Linux is useful for restoring and saving
#   OS-snapshots.
#
# Input:
#
#    1.) A tarball of a Ubuntu OS (we tested 16.04) root
#        partition. You can use the generic 'makeOsTgz.bash' script to
#        create this tarball.  Make sure to run makeOsTgz.bash
#        on the computer whose root partition you want to copy.
#
#    2.) Two optional tgz files: freedos.tgz and gparted.tgz.  They
#        should be present in the same directory as this script.
#
# Output:
#
#    1.) This script first creates the following disk paritions
#        (assuming the two harddrives are /dev/sda and /dev/sdb):
#
#       /dev/sda:
#          sda1: 512MB           ---- Free DOS
#          sda2: 20GB            ---- Your provided Linux OS goes here
#          sda3: 20GB            ---- Swap
#          sda4: Rest of Space   -----Linux raid autodetect
#
#       /dev/sdb:
#          sdb1: 512MB           ---- Gparted Linux OS
#          sdb2: 20GB            ---- Future OS
#          sdb3: 20GB            ---- OS Snapshot Storage
#          sdb4: Rest of Space   ---- Linux raid autodetect
#
#    2.) Installs Free DOS to /dev/sda1.
#    3.) Extracts the provided OS partition tarball (eg, rootDir.tgz) to
#        /dev/sda2, and sets up Grub and /etc/fstab accordingly (eg.,
#        adds Grub entries for swap and /share -- see below).
#    4.) Formats the swap partition in /dev/sda3.
#    5.) Installs Gparated Linux to /dev/sdb1.
#    6.) Formats /dev/sdb3 so that we can store OS snapshots (compressed).
#    7.) Creates a RAID0 partition combining /dev/sda4 and /dev/sdb4 and then
#        mounts this partition as /share.
#    8.) Creates a "symbolic link" from /var/lib/glance/images to
#       /share/OS_SymLinks/sda2/var/lib/glance/images.
#    9.) Creates a "symbolic link" from /var/lib/nova/{instances,images} to
#       /share/OS_SymLinks/sda2/var/lib/nova/{instances,images}.
#
#    Note that /dev/sdb2 is empty and can be used for installing
#    another OS.
#
# Log:
#   2017-12-16 (Sat) -- Provide alternative to /dev/sda and /dev/sdb
#   2017-04-07 (Fri) -- bind /var/lib/nova/instances and /var/lib/glance/images
#   2017-04-04 (Tue) -- Added Snapshot Grub menuitems
#   2017-04-03 (Mon) -- Initial public release
#   2017-03-30 (Thu) -- Initial version
#
#
set -e

MY_NAME=`basename $0`

function err_report {
   echo "Fail @ [${MY_NAME}:${1}]"
   exit 1
}

trap 'err_report $LINENO' ERR


read -r -d '' usageStr << EOM || true

Incorrect Usage:
  $0 $*

Correct Usage:
  $0 [options] <rootDirTgz>

Required:
  <rootDirTgz>  -- The OS partition tarball.

Options:
  -1 <devPath>  -- Drive 1 device. Default /dev/sda
  -2 <devPath>  -- Drive 2 device. Default /dev/sdb

Example:
  $0 rootDir.tgz 
  $0 rootDir.tgz -1 /dev/sdb -2 /dev/sdc
EOM


# default parameter
DRIVE1=/dev/sda
DRIVE2=/dev/sdb

# handle arguments:
nRequired=1
pIdx=0
while [[ $1 ]]; do
    case "$1" in
	(-h | --help)
	    echo "$usageStr"
	    exit 0;;
	(-1 | --drive1)
	    DRIVE1=$2
	    echo "DRIVE1 is now set to $DRIVE1"
	    shift 2;;
	(-2 | --drive2)
	    DRIVE2=$2
	    echo "DRIVE2 is now set to $DRIVE2"
	    shift 2;;
	(*)
	    case $pIdx in
		(0) 		# 1st positional argument
		    OSTGZ=$1
		    pIdx=$(($pIdx+1)); shift;;
		(*)
		    echo "$usageStr"
		    exit 1
	    esac
	    ;;
    esac
done
if [[ ! $pIdx -eq $nRequired ]]; then
    echo "$usageStr"
    exit 1;
fi


ROOTDIR_TGZ=$(readlink -f $OSTGZ)
FREEDOS_TGZ=$(readlink -f freedos.tgz)
GPARTED_TGZ=$(readlink -f gparted.tgz)
test -e $ROOTDIR_TGZ
test -e $FREEDOS_TGZ
test -e $GPARTED_TGZ


function mountOSAndChdir ()
{
    OS_MNTDIR=/mnt/os
    mkdir -p $OS_MNTDIR
    mount $osPart $OS_MNTDIR
    cd $OS_MNTDIR
}


# 1.) check utility and others
#--------------------------
echo "checking kpartx" && which kpartx
echo "checking ddrescue" && which ddrescue
echo "checking losetup" && which losetup
echo "checking rsync" && which rsync
echo "checking lsblk" && which lsblk
echo "checking sed" && which sed
echo "checking grub-install" && which grub-install
echo "checking update-grub2" && which update-grub2
echo "checking mkswap" && which mkswap
echo "checking mkfs.ext4" && which mkfs.ext4
echo "checking e2label" && which e2label
echo "checking fdisk" && which fdisk
echo "checking uuidgen" && which uuidgen
echo "checking sync" && which sync
echo "checking awk" && which awk
echo "checking cut" && which cut

if [[ ! $USER = "root" ]]; then
    echo "must be root or sudo"
    exit 1
fi

if [[ ! -e $ROOTDIR_TGZ ]]; then
    echo "$ROOTDIR_TGZ file not found"
    exit 1
fi


# 2.) Partition the pair of matching HDs
#--------------------------

## Check to makesure the pair of hd matches:
echo "Checking the Hardddrives"
#cat /proc/partitions | grep -q $DRIVE1
#cat /proc/partitions | grep -q $DRIVE2
fdisk -l $DRIVE1
fdisk -l $DRIVE2
DRIVE1_INFO=$(fdisk -l $DRIVE1 | head -n 1 | cut -d ' ' -f 3 )
DRIVE2_INFO=$(fdisk -l $DRIVE2 | head -n 1 | cut -d ' ' -f 3 )
test "$DRIVE1_INFO" == "$DRIVE2_INFO"
fdisk -l $DRIVE1
fdisk -l $DRIVE2

echo ""
echo ""
echo  "WARNING WARNING WARNING WARNING WARNING WARNING WARNING "
echo  "About to WIPE harddrives!! "
echo  "Please hit enter to contiue or hit ctrl-c to abort."
read 

## undo raid if necessary.
mdadm -S /dev/md127 || true
dd if=/dev/zero of=$DRIVE1 bs=512 count=10
dd if=/dev/zero of=$DRIVE2 bs=512 count=10

## Create the 4 partitions:
echo "Creating the Partition"
echo -e "o\nn\np\n1\n\n+512M\nn\np\n2\n\n+20G\nn\np\n3\n\n+20G\nn\np\n\n\nt\n1\nde\nt\n3\n82\nt\n4\nfd\nw\nq\n" | fdisk -u $DRIVE1
sleep 1
sync

echo -e "o\nn\np\n1\n\n+512M\nn\np\n2\n\n+20G\nn\np\n3\n\n+20G\nn\np\n\n\nt\n4\nfd\nw\nq\n" | fdisk -u $DRIVE2
sleep 1
sync

## Format the partitions:
echo "Formatting the Partitions"
echo "  Disk1 $DRIVE1:"
freedosGrub=(hd0,msdos1)
freedosPart=${DRIVE1}1
mkfs.vfat $freedosPart
osPart=${DRIVE1}2
osDrive=${DRIVE1}
mkfs.ext4 $osPart
swapPart=${DRIVE1}3
mkswap -U `uuidgen` $swapPart

echo "  Disk2 $DRIVE2:"
gpartedGrub=(hd1,msdos1)
gpartedPart=${DRIVE2}1
mkfs.ext4 $gpartedPart
mkfs.ext4 ${DRIVE2}2
mkfs.ext4 ${DRIVE2}3
e2label ${DRIVE2}3 SNAP_PART

echo "  RAID0:"
sharePart=/dev/md127
mdadm -C $sharePart --level=raid0 --raid-devices=2 ${DRIVE1}4 ${DRIVE2}4
mkfs.ext4 $sharePart



# 3.) Install the Provided OS
#--------------------------
## Transfer the providedd OS:
echo "Transfering the provided Ubuntu OS"
mountOSAndChdir
time tar -xf $ROOTDIR_TGZ

## setup the /etc/fstab
newRootUUID=`tune2fs -l $osPart | grep UUID | awk '{print \$3;}'`
echo $newRootUUID
echo "UUID=$newRootUUID  /   ext4    errors=remount-ro  0 1" > etc/fstab

newSwapUUID=`swaplabel $swapPart | awk '{print \$2;}'`
echo $newSwapUUID 
echo "UUID=$newSwapUUID  none swap    sw                0 0" >> etc/fstab

newShareUUID=`tune2fs -l $sharePart | grep UUID | awk '{print \$3;}'`
echo $newShareUUID
echo "UUID=$newShareUUID /share ext4 defaults           0 2" >> etc/fstab

echo "/share/OS_SymLinks/$osPart/var/lib/glance/images \
/var/lib/glance/images none bind"  >> etc/fstab
mkdir -p share/OS_SymLinks/$osPart/var/lib/glance/images || true
chmod a+rw share/OS_SymLinks/$osPart/var/lib/glance/images || true
mkdir -p var/lib/glance/images || true

echo "/share/OS_SymLinks/$osPart/var/lib/nova/images \
/var/lib/nova/images none bind" >> etc/fstab
echo "/share/OS_SymLinks/$osPart/var/lib/nova/instances \
/var/lib/nova/instances none bind" >> etc/fstab
mkdir -p share/OS_SymLinks/$osPart/var/lib/nova/{images,instances} || true
chmod a+rw share/OS_SymLinks/$osPart/var/lib/nova/{images,instances} || true
mkdir -p var/lib/nova/{images,instances} || true

## clean up
cd; umount $OS_MNTDIR


# 3.) Prepare Custome Grub entries
#--------------------------
mountOSAndChdir
cat <<"EOF" > etc/grub.d/40_custom
#!/bin/sh
exec tail -n +3 $0

######
EOF
chmod a+x etc/grub.d/40_custom

## clean up
cd; umount $OS_MNTDIR



# 4.) Install Free DOS
#--------------------------
if [[ -e $FREEDOS_TGZ ]]; then
    echo "Installing Free DOS"
    MNTDIR=/mnt/freedos
    mkdir -p $MNTDIR
    mount $freedosPart $MNTDIR
    time tar -C $MNTDIR -xf $FREEDOS_TGZ
    umount $MNTDIR

    mountOSAndChdir
    cat <<EOF >> etc/grub.d/40_custom
menuentry 'FreeDos 1.0' {
    insmod fat
    set root='$freedosGrub'
    linux16 /freedos/memdisk
    initrd16 /freedos/freedos.img
}

EOF
    ## clean up
    cd; umount $OS_MNTDIR
else
    echo "No freedos.tgz found, skip installing Free DOS"
fi


# 5.) Install Gparted Linux
#--------------------------
if [[ -e $GPARTED_TGZ ]]; then
    echo "Installing Gparted Linux"
    MNTDIR=/mnt/gparted
    mkdir -p $MNTDIR
    mount $gpartedPart $MNTDIR
    time tar -C $MNTDIR -xf $GPARTED_TGZ
    umount $MNTDIR
    read -p "Are you booting from USB? (y/n)?"
    if [ $REPLY != "y" ]; then
	echo "setting gpartedPart to /dev/sdb1 for grub entry"
	gpartedPart=/dev/sdb1
    fi

    mountOSAndChdir
    cat <<EOF >> etc/grub.d/40_custom
menuentry 'GParted live' {
    set root='$gpartedGrub'
    linux /gparted/live/vmlinuz boot=live config union=aufs noswap noprompt nomodeset vga=normal ip=frommedia live-media-path=/gparted/live bootfrom=$gpartedPart toram=filesystem.squashfs
    initrd /gparted/live/initrd.img
}

menuentry 'Take OS Snapshot' {
    set root='$gpartedGrub'
    linux /gparted/live/vmlinuz boot=live config union=aufs noswap noprompt nomodeset vga=normal ip=frommedia live-media-path=/gparted/live bootfrom=$gpartedPart toram=filesystem.squashfs takeOsSnapshot gl-debug 2
    initrd /gparted/live/initrd.img
}


menuentry 'Restore OS Snapshot' {
    set root='$gpartedGrub'
    linux /gparted/live/vmlinuz boot=live config union=aufs noswap noprompt nomodeset vga=normal ip=frommedia live-media-path=/gparted/live bootfrom=$gpartedPart toram=filesystem.squashfs restoreOsSnapshot gl-debug 2
    initrd /gparted/live/initrd.img
}


EOF
    ## clean up
    cd; umount $OS_MNTDIR
else
    echo "No gparted.tgz found, skip installing Gparted Linux"
fi


# 6.) Install Grub Last
#--------------------------
## Install GRUB
echo "Installing GRUB"
mountOSAndChdir
for i in /dev /dev/pts /proc /sys /run; do
    mkdir -p $OS_MNTDIR$i; mount -B $i $OS_MNTDIR$i;
done
cat <<"EOF" > $OS_MNTDIR/makeBootable.bash
#!/bin/bash
cd /
grub-install --root-directory=/ --no-floppy --modules="normal part_msdos ext2 multiboot biosdisk" $osDrive
echo -e "\nGRUB_DISABLE_OS_PROBER=false\n" >> etc/default/grub
update-grub2
EOF
chmod a+x $OS_MNTDIR/makeBootable.bash
chroot $OS_MNTDIR env osDrive=$osDrive ./makeBootable.bash

## Clean ups:
for i in /dev/pts /dev /proc /sys /run; do umount -f $OS_MNTDIR$i; done
cd; umount $OS_MNTDIR


# 7.) Done 
#--------------------------
echo ""
echo ""
echo "Done!  You can reboot now."