#!/bin/bash -e
#
# Authors:
#   Tommy Chang (chang177@purdue.edu)
#
# Version 2019-10-05 (Sat)
#
# Description:
#   Run this script in an empty directory to create a VM of the Linux
#   OS that is currently running on your computer.  This script has
#   been tested on Ubuntu 14.04, 16.04, and 18.04.
#
#   After running this script, you will find a blockXXX.dd file in the
#   current directory.  Simply upload this VM file to your cloud.
#
# Log:
#   2019-10-06 (Sat) -- Exclude sparse file /var/log/lastlog
#   2019-02-09 (Sat) -- Fix swap error.
#   2019-01-11 (Fri) -- Allow root filesystem to be other than /
#   2018-12-04 (Tue) -- use swap file instead of partition (untested)
#   2018-06-15 (Fri) -- tested
#   2017-06-16 (Fri) -- fixed home directory permission bug (untested)
#   2017-05-28 (Sun) -- added 3rd option for copying home directory
#   2017-05-04 (Thu) -- Copy minimal /home/
#   2017-04-03 (Mon) -- Initial public release.
#   2017-03-30 (Thu) -- perl is not needed, more checks and comments
#   2017-03-27 (Mon) -- add excludeDir_flg
#   2017-03-01 (Wed) -- overhaul and added comments
#   2016-11-23 (Wed) -- Initially created 
#
set -e

RUN_FROM_SCRATCH="Yes"
# debug
#set -x
#RUN_FROM_SCRATCH="No"
#MOUNT_ONLY="Yes"


if [[ ! ($# -eq 4)]]; then
    cat <<EOF
Usage:
  $0 <DiskSizeInGB> <SwapSizeInGB> <rootFileSys> <copyHomeFlag>

Example:
  # Create a VM with 15GB total space and 0.5GB for swap.
  # Also, create a minimal home directory (don't copy the actual /home)
  sudo $0 15 0.5 / 0
  # Same as the previous example, except the fileystem is mounted on /mnt/fs
  sudo $0 15 0.5 /mnt/fs/ 0
EOF
    exit 1;
fi


echo "=================================="
echo "# 1.) check utility"
echo "=================================="
which kpartx > /dev/null || (echo "could not find kpartx" && false)
which ddrescue > /dev/null || (echo "could not find ddrescue" && false)
which losetup > /dev/null || (echo "could not find losetup" && false)
which rsync > /dev/null || (echo "could not find rsync" && false)
which lsblk > /dev/null || (echo "could not find lsblk" && false)
which sed > /dev/null || (echo "could not find sed" && false)
which grub-install > /dev/null || (echo "could not find grub-install" && false)
which update-grub2 > /dev/null || (echo "could not find update-grub2" && false)
which mkswap > /dev/null || (echo "could not find mkswap" && false)
which mkfs.ext4 > /dev/null || (echo "could not find mkfs.ext4" && false)
which fdisk > /dev/null || (echo "could not find fdisk" && false)
which uuidgen > /dev/null || (echo "could not find uuidgen" && false)
which find > /dev/null || (echo "could not find find" && false)
which wc > /dev/null || (echo "could not find wc" && false)
which fallocate > /dev/null || (echo "could not find fallocate" && false)

[ -e /usr/lib/grub/i386-pc/modinfo.sh ] || (echo "can not find modinfo.sh, run: apt-get install grub-pc-bin" && false)

echo "Last chance to abort.  Do not interrupt.."
read -p "Press Enter to continue or control-c to abort NOW."

function cleanupLoopDev {
    echo "Error/Abort detected!...."
    echo "PLEASE WAIT while cleaning  up...."
    echo "umount $MNTDIR"
    sync
    umount $MNTDIR 2> /dev/null; sleep 1
    echo "rmdir $MNTDIR"
    rmdir $MNTDIR
    echo "losetup -d $loop1"
    losetup -d $loop1; sleep 1
    echo "kpartx -d  $loop0"
    kpartx -d  $loop0; sleep 1
    echo "losetup -d  $loop0"
    losetup -d  $loop0
    echo "Please run 'mount' and 'lsblk' to make sure everything is OK.."
    exit 1
}

trap cleanupLoopDev ERR
DISKSIZE=$1
SWAPSIZE=$2
ROOTDIR=$3
COPYHOME=$4

EXCLUDE_DIR=`pwd -P`
numSubDirs=$(find $EXCLUDE_DIR -type d | wc -l)
if [[ ! $numSubDirs -eq 1 ]]; then
    echo ""
    echo  "Please run this script from an empty directory."
    echo  "The stuff in the current directory will"
    echo  "not be copied to the VM."
    exit 1
fi


BLOCKDD="block${DISKSIZE}GB.dd"
MNTDIR=/mnt/junk$$

if [[ ! $USER = "root" ]]; then
    echo ""
    echo "Must run this script as root."
    exit 1
fi

if [[ "$RUN_FROM_SCRATCH" == "Yes" ]]; then
    echo "=================================="
    echo "# 2.) Create an empty dd file"
    echo "=================================="
    ddrescue -t -n -S -s${DISKSIZE}000Mi /dev/zero $BLOCKDD
    echo -e "o\nn\np\n1\n\n\nw\nq\n" | fdisk -u $BLOCKDD
    fdisk -lu $BLOCKDD
fi

echo "=================================="
echo "# 3.) Mount the dd disk"
echo "=================================="
echo "# whole disk"
loop0=`losetup -f`
losetup $loop0 $BLOCKDD
kpartx -av $loop0
sleep 1
echo "# mount 1st partition = root"
loop0p1=/dev/mapper/`lsblk -l $loop0 | grep p1 | awk '{print $1;}' | grep 'p1$'`; echo $loop0p1
loop1=`losetup -f`
losetup $loop1 $loop0p1

if [[ "$RUN_FROM_SCRATCH" == "Yes" ]]; then
    echo "=================================="
    echo "# 4.) format the first partition (in the dd file)"
    echo "=================================="
    mkfs.ext4 $loop1
fi

echo "=================================="
echo "# 5.) Install the OS  (just copy / to the first partition)"
echo "=================================="
echo " Please be patient"
mkdir -p $MNTDIR || true
mount $loop1 $MNTDIR

if [[ "$MOUNT_ONLY" == "Yes" ]]; then
    echo "mounted"
    exit 0
fi

# EXCLUDE_OPTION="--exclude var/lib/lightdm/.gvfs"
EXCLUDE_OPTION='--exclude */.gvfs'
EXCLUDE_OPTION="$EXCLUDE_OPTION --exclude ${ROOTDIR}var/log/lastlog"
if [[ "$COPYHOME" == "0" ]]; then
    EXCLUDE_OPTION="$EXCLUDE_OPTION --exclude ${ROOTDIR}home/"
fi
## exclude the swap file from being copied, if any
SWAP_FILE=`swapon -s |  grep "file" | awk '{print $1;}'`
if [[ -n $SWAP_FILE ]]; then
    EXCLUDE_OPTION="$EXCLUDE_OPTION --exclude $SWAP_FILE"
fi

if [[ "$RUN_FROM_SCRATCH" == "Yes" ]]; then
    rsync --info=progress2 -aHx $EXCLUDE_OPTION --exclude $EXCLUDE_DIR $ROOTDIR $MNTDIR 
    echo "# do it again in case something changed in ${ROOTDIR}"
    sync; sync;
    rsync --info=progress2 -aHx $EXCLUDE_OPTION --exclude $EXCLUDE_DIR $ROOTDIR $MNTDIR 
fi

if [[ "$COPYHOME" == "0" ]]; then
    echo "=================================="
    echo "# 6.) Copy the minimal home "
    echo "=================================="
    find ${ROOTDIR}home/ -maxdepth 2 -type f -iname ".*" | xargs tar -cvf home.tar
    sudo tar -C $MNTDIR -xvf home.tar
    rm home.tar

    # TC 2017-06-16 (Fri) Now we need to fix the permission: 
    allUserHomes=$(find $MNTDIR/home/* -maxdepth 0 -type d)
    for userHome in $allUserHomes; do
	aFile=$(find $userHome -maxdepth 2 -type f -iname ".*" | head -1)
	uid=$(stat -c%u $aFile)
	gid=$(stat -c%g $aFile)
	sudo chown $uid $userHome
	sudo chgrp $gid $userHome
    done
fi


echo "=================================="
echo "# 7.) Make dd (ie., the whole disk) bootable"
echo "=================================="
for i in /dev /dev/pts /proc /sys /run; do mkdir -p $MNTDIR$i; mount -B $i $MNTDIR$i; done

cat <<"EOF" > $MNTDIR/makeBootable.bash
#!/bin/bash
cd /
grub-install --root-directory=/ --no-floppy --modules="normal part_msdos ext2 multiboot biosdisk" $loop0
echo -e "\nGRUB_DISABLE_OS_PROBER=true\n" >> etc/default/grub
update-grub2
sed --in-place '/loop/d' boot/grub/grub.cfg
sed --in-place 's/quiet splash//g' boot/grub/grub.cfg
newRootUUID=`tune2fs -l $loop1 | grep UUID | awk '{print \$3;}'`; echo $newRootUUID

echo "UUID=$newRootUUID  /   ext4    errors=remount-ro 0 1" > etc/fstab
if ! [[ $SWAP_MB -eq 0 ]]; then
  fallocate -l ${SWAP_MB}M /swapfile
  mkswap /swapfile
  chmod 0600 /swapfile
  echo "/swapfile              none    swap sw           0 0" >> etc/fstab
fi

if [[ -e etc/udev/rules.d/70-persistent-net.rules ]]; then
   rm etc/udev/rules.d/70-persistent-net.rules
fi
EOF
chmod a+x $MNTDIR/makeBootable.bash
SWAP_MB=`echo "$SWAPSIZE * 1000" | bc -l`
SWAP_MB=`echo $SWAP_MB|cut -f1 -d"."`
echo "Round SWAP to $SWAP_MB MB"
chroot $MNTDIR env loop0=$loop0 loop1=$loop1 SWAP_MB=$SWAP_MB ./makeBootable.bash

echo "=================================="
echo "# 8.) Cleanup and unmount dd"
echo "=================================="
for i in /dev/pts /dev /proc /sys /run; do umount -f $MNTDIR$i; done
echo "Check to space below: "
df -h | grep $MNTDIR
cd;  umount $MNTDIR; rmdir $MNTDIR; sleep 1
echo "losetup -d $loop1"
losetup -d $loop1; sleep 1
echo "kpartx -d  $loop0"
kpartx -d  $loop0; sleep 1
echo "losetup -d  $loop0"
losetup -d  $loop0

echo "Checking to see of unmounted properly:"
test -z $(losetup -a | grep  $loop0)
echo "Done!  Your VM is $BLOCKDD"