#!/bin/bash

# Create a swap partition of the same size as the amount of physical RAM
# and a /home partition using all remaining space on the drive
# ronald.mccollam@canonical.com

DEVICE=sda
ROOTSIZE=20000

# Find the total physical RAM in the system
# (/proc/meminfo stores this as kB, but parted expects MB, so convert as well)
PHYSICALRAM=`grep MemTotal /proc/meminfo | awk '{ print $2 }'`
MEMSIZE=`echo $PHYSICALRAM/1024 | bc`

if [ "$MEMSIZE" -eq 0 ]
then
	echo "Unable to determine physical RAM size!"
	exit 1
fi
# Because we have the total 'available' physical RAM, we need some padding.
# We bump this by 1GB to be sure we have enough space to hibernate.
let MEMSIZE="$MEMSIZE"+1024



while getopts ":d:s" opt
do
	case $opt in
		d)
			DEVICE=$OPTARG
			;;
		\?)
			echo "Create a swap partition of the same size as system RAM, then resize the"
			echo "first partition to take all available space."
			echo
			echo "WARNING: Use only on a new system, as this script may delete existing partitions!"
			echo
			echo "Usage: $0 [-d device] [-s size]"
			echo "-d: Specifies device to read (default: sda)"
			echo "-s: Specifies the size of the root partition in MB (default: $ROOTSIZE)"
			exit 255
			;;
	esac
done

if [ $UID -ne 0 ]
then
	echo "You must have root permissions to run this script."
	exit 2
fi

##### Sanity-check the environment

# we expect the partitions to have been created as per the disk recipe
# on the image - 20G /, plus an extended partition occupinging the
# rest.  Check they're not mounted and destroy them.
MOUNTCOUNT=`fdisk -l /dev/"$DEVICE" | grep "^/dev/$DEVICE" | wc -l`
if [ $MOUNTCOUNT -eq 3 ]
then
    for DEV in $(fdisk -l /dev/"$DEVICE" | grep "^/dev/${DEVICE}[2-9]"); do
        if mount | grep "^$DEV"; then
	    echo "Partition $DEV found but mounted.  Unrecognized configuration"
	    exit 3
        fi
    done
    # there are holes in the above checking, but if you're running
    # this and created your own partitions, you are on your own.

    # only have to remove the extended partition (2):
    parted --script -- /dev/$DEVICE rm 2
fi

# Make sure only one partition exists
MOUNTCOUNT=`fdisk -l /dev/"$DEVICE" | grep "^/dev/$DEVICE" | wc -l`
if [ $MOUNTCOUNT -gt 1 ]
then
	echo "More than one partition already exists.  Aborting."
	exit 3
fi

# Make sure there is empty space at the end of the disk
CYLINDERS=`fdisk -l /dev/"$DEVICE" | grep cylinders\$ | awk '{ print $5 }'`
USED=`fdisk -l /dev/"$DEVICE" | grep "$DEVICE"1 | sed s/\*// | awk '{ print $3 }'`
if [ $USED -ge $CYLINDERS ]
then
	echo "No empty space at end of disk.  Aborting."
	exit 4
fi

# Make sure no swap space is allocated/in use
SWAPTOTAL=`free | grep Swap | awk '{ print $2 }'`
if [ $SWAPTOTAL -ne 0 ]
then
	echo "Swap exists and is in use.  Aborting."
	exit 5
fi


##### Create a swap partition
# It should start at the end of the root partition and end $MEMSIZE later...
let MEMEND="$ROOTSIZE"+"$MEMSIZE"
ROOTSTRING="$ROOTSIZE"MB
MEMSTRING="$MEMEND"MB
parted --script -- /dev/$DEVICE mkpart primary linux-swap $ROOTSTRING $MEMSTRING
if [ $? -ne 0 ]
then
	echo "CRITICAL FAILURE: Creating swap partition failed!"
	exit 1
fi

sync
mkswap /dev/"$DEVICE"2
sync

##### Use the remaining space for the data partition
# It should start at the end of the swap partition and end at the end of the disk
# ("-0" is the end of the disk)
parted --script -- /dev/$DEVICE mkpart primary ext4 $MEMSTRING -0
if [ $? -ne 0 ]
then
	echo "CRITICAL FAILURE: Creating /home partition failed!"
	exit 1
fi

sync
mkfs.ext4 /dev/"$DEVICE"3
sync

##### Mount everything
# Add the new partitions to fstab:
echo "/dev/$DEVICE"2	swap	swap	defaults	0	0 >> /etc/fstab
echo "/dev/$DEVICE"3	/home	ext4	defaults	0	0 >> /etc/fstab
# Grab the OOBE user data (this way we can still run through the initial setup
#after blowing away the old /home)
mv /home/oem /tmp
mount -a
swapon /dev/"$DEVICE"2
sync
mv /tmp/oem /home
