export TEMP_MOUNT=$(mktemp -d /tmp/TEMP.XXXXX)
export DISK_ROOT_MOUNT=$(mktemp -d /tmp/DISKROOT.XXXXX)
export DISK_BOOT_MOUNT=$(mktemp -d /tmp/DISKBOOT.XXXXX)
export IMAGE_ROOT_MOUNT=$(mktemp -d /tmp/IMAGE.XXXXX)
export DD_FILE=${TEMP_MOUNT}/image.dd
BTRFS_MOUNT_OPTS="noatime,compress,nossd,skip_balance,nospace_cache"
BTRFS_MOUNT_OPTS_SCRATCH="noatime,compress-force=zlib,nossd,skip_balance,nospace_cache"
export DEBUG=${DEBUG:-0}
export IMAGE_LOG=${IMAGE_LOG:-/dev/kmsg}
export FLASH_DEVICE=${FLASH_DEVICE:-$FLASH_DRIVE}
 
getDiskInfo(){
	CAPTURE_BOOT_START=$(parted $FLASH_DEVICE unit KB print | grep "^\ 1 " | awk '{print $2}')
	CAPTURE_BOOT_END=$(parted $FLASH_DEVICE unit KB print | grep "^\ 1 " | awk '{print $3}')
	CAPTURE_BOOT_SIZE_KB=$(parted $FLASH_DEVICE unit KB print | grep "^\ 1 " | awk '{print $4}' | grep -o "[0-9]*")
	
	# Root Partition Info
	CAPTURE_ROOT_START=$(parted $FLASH_DEVICE unit KB print | grep "^\ 2 " | awk '{print $2}')
	CAPTURE_ROOT_END=$(parted $FLASH_DEVICE unit KB print | grep "^\ 2 " | awk '{print $3}')
	CAPTURE_ROOT_PART_SZ_ORIG_KB=$(parted $FLASH_DEVICE unit KB print | grep "^\ 2 " | awk '{print $4}' | grep -o "[0-9]*")
	CAPTURE_ROOT_SIZE_KB=$(( $(getBtrfsUsage) * 1024))

	CAPTURE_ROOT_UUID=`blkid -o udev $ROOT_PART | grep ID_FS_UUID= | cut -d= -f2`
	CAPTURE_ROOT_NODE_SZ=`cat /sys/fs/btrfs/$CAPTURE_ROOT_UUID/nodesize 2>/dev/null`
	CAPTURE_ROOT_FEATURE_EXTENDED_IREF=`cat /sys/fs/btrfs/$CAPTURE_ROOT_UUID/features/extended_iref 2>/dev/null`
	CAPTURE_ROOT_FEATURE_EXTENDED_IREF=${CAPTURE_ROOT_FEATURE_EXTENDED_IREF:-0}
	CAPTURE_ROOT_FEATURE_MIXED_BACKREF=`cat /sys/fs/btrfs/$CAPTURE_ROOT_UUID/features/mixed_backref 2>/dev/null`
	CAPTURE_ROOT_FEATURE_MIXED_BACKREF=${CAPTURE_ROOT_FEATURE_MIXED_BACKREF:-0}
	CAPTURE_ROOT_FEATURE_SKINNY_METADATA=`cat /sys/fs/btrfs/$CAPTURE_ROOT_UUID/features/skinny_metadata 2>/dev/null`
	CAPTURE_ROOT_FEATURE_SKINNY_METADATA=${CAPTURE_ROOT_FEATURE_SKINNY_METADATA:-0}
	CAPTURE_ROOT_FEATURE_NO_HOLES=`cat /sys/fs/btrfs/$CAPTURE_ROOT_UUID/features/no_holes 2>/dev/null`
	CAPTURE_ROOT_FEATURE_NO_HOLES=${CAPTURE_ROOT_FEATURE_NO_HOLES:-0}

	MKFS_BTRFS_OPTIONS="-f -n ${CAPTURE_ROOT_NODE_SZ} -O "
	if [ $CAPTURE_ROOT_FEATURE_EXTENDED_IREF -eq 1 ] ; then
		MKFS_BTRFS_OPTIONS+="extref,"
	else
		MKFS_BTRFS_OPTIONS+="^extref,"
	fi
	if [ $CAPTURE_ROOT_FEATURE_SKINNY_METADATA -eq 1 ] ; then
		MKFS_BTRFS_OPTIONS+="skinny-metadata,"
	else
		MKFS_BTRFS_OPTIONS+=^"skinny-metadata,"
	fi
	if [ $CAPTURE_ROOT_FEATURE_NO_HOLES -eq 1 ] ; then
		MKFS_BTRFS_OPTIONS+="no-holes"
	else
		MKFS_BTRFS_OPTIONS+="^no-holes"
	fi

	# Flag info
	BOOT_FLAG_PT="$(parted $FLASH_DEVICE print | grep boot | awk '{print $1}')"
	BOOT_PART_TYPE="$(parted $FLASH_DEVICE print | grep "^ 1 " | awk '{print $6}')"
	LBA_FLAG_PT="$(parted $FLASH_DEVICE print | grep lba | awk '{print $1}')"
}
 
cleanUp() {
	[ $DEBUG -ne 0 ] && echo "Cleaning up and unmounting filesystems"

	[ -n "`mount | grep $IMAGE_ROOT_MOUNT`" ] && umount $IMAGE_ROOT_MOUNT

	# Cleanup DD (unmount and remove loopbacks
	LODEVS=$(losetup -a | cut -f1 -d " ")
	for lodev in $LODEVS ; do
		losetup -d ${lodev::-1}
	done

	sleep 1
	umount -l $TEMP_MOUNT
	
	# Remove all of the read only subvolumes
	if [ -n "`mount | grep $DISK_ROOT_MOUNT`" ]
	then
		btrfs subvol list $DISK_ROOT_MOUNT | grep _ro | cut -d ' ' -f9 | xargs -I {} btrfs subvol del $DISK_ROOT_MOUNT/{}  >>$IMAGE_LOG 2>&1
		sync
		
		[ -f $DISK_ROOT_MOUNT ] && rm $DISK_ROOT_MOUNT/origroot_boundaries >> $IMAGE_LOG 2>&1
		umount $DISK_ROOT_MOUNT
	fi

	[ -n "$TEMP_PART" ] && delete_partition $TEMP_PART
	
	setupPostProcessing
}

setupPostProcessing() {
	mkdir -p /tmp/tmp_boot 2>/dev/null
	mount $BOOT_PART /tmp/tmp_boot

	# from thinstate
	[ -e /usr/share/postimg/hooks.delayed/10_resize_root ] && cp /usr/share/postimg/hooks.delayed/10_resize_root /tmp/tmp_boot/POSTIMG.DIR/DELAYED/
	# from ukit
	[ -e $USBKEY/hooks/10_resize_root ] && cp $USBKEY/hooks/10_resize_root /tmp/tmp_boot/POSTIMG.DIR/DELAYED/

	chmod +x /tmp/tmp_boot/POSTIMG.DIR/DELAYED/10_resize_root

	umount /tmp/tmp_boot	
}

getDiskByLabel() {
	local label=$1
	blockdev --rereadpt ${FLASH_DEVICE} >> $IMAGE_LOG  2>&1
	local disk=`blkid ${FLASH_DEVICE}* | grep LABEL=\"${label}\" | cut -d: -f1`
	
	if [ "x$disk" == "x" ]; then 
		disk=`blkid ${FLASH_DEVICE}* | grep LABEL=\"$2\" | cut -d: -f1`
	fi

	echo $disk
}

getDiskSizeKb() {
	local s=`blockdev --getsize64 $1`
	echo $((s/1024))
}

getBtrfsUsage(){
	[ $DEBUG -ne 0 ] && set -x

	# probe ROOT to see if clone targe is a 32 or 64 bit system, setting floor accordingly
	[ -d ${DISK_ROOT_MOUNT}/root/lib/ ] && FLOOR=952
	[ -d ${DISK_ROOT_MOUNT}/root/lib64/ ] && FLOOR=1900

	dataUsedMB=$(returnMB $(btrfs fi df $DISK_ROOT_MOUNT | grep Data | grep -o used=[0-9a-zA-Z.]* | cut -d '=' -f2))
	sysTotalMB=$(returnMB $(btrfs fi df $DISK_ROOT_MOUNT | grep System | grep -o used=[0-9a-zA-Z.]* | cut -d '=' -f2 | cut -d ',' -f1))
	metaUsedMB=$(returnMB $(btrfs fi df $DISK_ROOT_MOUNT | grep Meta | grep -o used=[0-9a-zA-Z.]* | cut -d '=' -f2))
	
	SZ=`printf "%.0f" "$(echo "($dataUsedMB + $sysTotalMB + $metaUsedMB )" | bc -l)"`
	
	# Enforce a floor of $FLOOR for the root file system if it is less than 952
	if [ $SZ -le $FLOOR ]
	then
		echo $FLOOR
	else
		printf "%.0f" "$(echo "($SZ * 1.3)" | bc -l)"
	fi
}

getMinimumBTRFSPartSize() {
	# we want the result to be a percentage larger than current usage.  Set that percentage here.
	local STEP_PCT=1.5
	local size=""

	# Minimum size set to a percentage over current usage (data + metadata)
	size=$(echo "$(getBtrfsUsage)*$STEP_PCT" | bc)
	echo $size
}

returnMB() {
 	local size=${1//[i,]/}
	
	local notation=${size: -2:2} # Find out whether or not we're using GB or MB (it varies depending on disk usage) and
	local newsize=

	# If given size is in GB, convert to MB (dropping the GB along the way)
	if [ x$notation = xGB ] ; then
		newsize=$(echo "${size%%GB} * 1024" | bc)
		newsize=${newsize%%\.*}
	elif [ x$notation = xKB ] ; then
		newsize=$(echo "${size%%KB} / 1024" | bc)
		newsize=${newsize%%\.*}
	else # Otherwise, just drop the MB and the decimal as we'll be adding a buffer later on anyway
		newsize=${size%%\.*}
	fi

	echo $newsize
}

create_temp_partition() {
	[ $DEBUG -ne 0 ] && set -x
	
	local device=$1
	local new_partnum=$(( $(parted $FLASH_DEVICE unit KB print | grep "^\ [0-9]" | tail -1 | awk '{print $1}') + 1 ))
	local new_partstart=$(parted $FLASH_DEVICE unit KB print | grep "^\ [0-9]" | tail -1 | awk '{print $3}')

        ls ${FLASH_DEVICE}* | sort > /tmp/before.txt

	if [ "x$tableType" = "xMBR" ] ; then
		MKPARTOPT="primary ext2"
	else 
		MKPARTOPT="TEMP ext2"
	fi
	parted -s $FLASH_DEVICE mkpart $MKPARTOPT $new_partstart ${DISK_SIZE_KB}kB >> $IMAGE_LOG 2>&1

	# Return the new partition name,
        ls ${FLASH_DEVICE}* | sort > /tmp/after.txt
        newpart=`diff /tmp/before.txt /tmp/after.txt |grep "^>" |awk '{print $2}'`
        if [ -b ${device}${new_partnum} ] ; then
                echo "${device}${new_partnum}"
        elif [ -b ${newpart} ] ; then
                echo "${newpart}"
        else
                echo "badpartname"
        fi
}

delete_partition() {
	[ $DEBUG -ne 0 ] && set -x
	local partition=$1
	local partnum=${partition: -1:1}

	parted -s $FLASH_DEVICE rm $partnum >> $IMAGE_LOG 2>&1
}

resize_partition(){
	[ $DEBUG -ne 0 ] && set -x
	local_disk=${ROOT_PART:: -1}
	local partnum=${ROOT_PART: -1}
	local offset=$1

	parted -s $FLASH_DEVICE rm 2 >> $IMAGE_LOG 2>&1
	if [ "x$tableType" = "xMBR" ] ; then
		MKPARTOPT="primary btrfs"
	else 
		MKPARTOPT="ROOT btrfs"
	fi
	parted -s $FLASH_DEVICE mkpart $MKPARTOPT $CAPTURE_ROOT_START $(( $CAPTURE_ROOT_PART_SZ_ORIG_KB - $offset))kB >> $IMAGE_LOG 2>&1
}

partition_dd(){
	[ $DEBUG -ne 0 ] && set -x
	local dd_file=$1

	# rebuild partition table
        if [ "x$tableType" = "xMBR" ] ; then
		parted -s $dd_file mktable msdos >> $IMAGE_LOG 2>&1
		parted -s $dd_file mkpart primary $BOOT_PART_TYPE $CAPTURE_BOOT_START $CAPTURE_BOOT_END >> $IMAGE_LOG 2>&1
		parted -s $dd_file mkpart primary btrfs $CAPTURE_BOOT_END 100%
        else
		parted -s $dd_file mktable gpt >> $IMAGE_LOG 2>&1
		parted -s $dd_file mkpart BOOT $BOOT_PART_TYPE $CAPTURE_BOOT_START $CAPTURE_BOOT_END >> $IMAGE_LOG 2>&1
		parted -s $dd_file mkpart ROOT btrfs $CAPTURE_BOOT_END 100%
	fi

	# Set bootable flag
	[ -n "$BOOT_FLAG_PT" ] && parted -s $DD_FILE set $BOOT_FLAG_PT boot on

	# Set LBA flag if applicable, notice that we check if the flag is already set then we don't do it again
	[[ -n "$LBA_FLAG_PT" && -n `parted -s $DD_FILE print | grep "^ $LBA_FLAG_PT " | grep -v lba` ]] && parted -s $DD_FILE set $LBA_FLAG_PT lba on
}

zeroFillFS() {
	[ $DEBUG -ne 0 ] && set -x
	local fs=$1
	
	sfill -Ifllz $fs
	btrfs fi sync $fs > /dev/null
	btrfs fi balance $fs
}

removeToollessBIOSFiles() {
	[ $DEBUG -ne 0 ] && set -x
	# according to toolless bios spec, we will try to clean up \HP\BIOS\Previous\ and zero out the unused sectors in BOOT partition

	# return if our BOOT is not vfat
	blkid $BOOT_PART | grep -q 'TYPE="vfat"'  || return

	mount $BOOT_PART $DISK_BOOT_MOUNT

	rm -rf $DISK_BOOT_MOUNT/EFI/HP/BIOSUpdate/Previous/* 2>/dev/null
	dd if=/dev/zero of=$DISK_BOOT_MOUNT/zeroFile >/dev/null  2>&1
	rm -rf $DISK_BOOT_MOUNT/zeroFile 2>/dev/null
	sync

	umount $DISK_BOOT_MOUNT
}

checkFreeDiskSpace() {
	[ $DEBUG -ne 0 ] && set -x
 	BOOT_PART=$(getDiskByLabel "BOOT" "HP_TOOLS")
	ROOT_PART=$(getDiskByLabel "ROOT")

	removeToollessBIOSFiles

	# clear away any hangover _ro snapshots
	mount -o noatime,compress,nossd,skip_balance,nospace_cache $ROOT_PART $DISK_ROOT_MOUNT
	for i in `btrfs subvol list -r $DISK_ROOT_MOUNT | awk '{print $9}' | grep _ro`; do
		btrfs subvol delete $DISK_ROOT_MOUNT/$i;
	done;
	
	# If a balance was in progress, let's complete it. Not completing the balance can result in an inflated Root partition usage
	# we skip the check now because btrfs-tools changes the return code of "btrfs balance status", which leaves us in the cold
	#  resume a non-existing balance job will return immediately so it's simple just to call resume without error-prone checks
	btrfs balance resume $DISK_ROOT_MOUNT >>$IMAGE_LOG 2>&1
			
	# Populate disk variables we need throughout the script
	getDiskInfo

	DD_SIZE_KB=$(( $CAPTURE_BOOT_SIZE_KB + $CAPTURE_ROOT_SIZE_KB + 50))

	DISK_SIZE_KB=$(parted $FLASH_DEVICE unit KB print | grep $FLASH_DEVICE | awk '{print $3}' | grep -o "[0-9]*")
	TEMP_PART_SIZE_KB=$(( $DD_SIZE_KB + 102400))

	DISK_FREE_SPACE_KB=$(( $DISK_SIZE_KB - $CAPTURE_ROOT_SIZE_KB - $CAPTURE_BOOT_SIZE_KB))
	

	# If we don't have enough space on the disk to use (total disk size minus the temp partition) we have to try to resize
	if [ $(( $(getDiskSizeKb $FLASH_DEVICE) - $(getDiskSizeKb $ROOT_PART) - CAPTURE_BOOT_SIZE_KB )) -le $TEMP_PART_SIZE_KB ] ; then
		if [ $(( DISK_FREE_SPACE_KB - TEMP_PART_SIZE_KB )) -gt 0 ] ; then
			btrfs fi resize -${TEMP_PART_SIZE_KB} $DISK_ROOT_MOUNT >> $IMAGE_LOG 2>&1
			if [ $? -ne 0 ] ; then
				echo "Error while resizing FS.  Exiting"
				RESULT=4
				exit 4
			fi

			btrfs fi sync $DISK_ROOT_MOUNT >> $IMAGE_LOG 2>&1
			
			# Before umounting the disk, place a reference file of the original root partition end point in the disk.
			# This allows us to restore the partition back to its original size in the event of abnormal termination of the capture process 
			# once we have created the scratch partition
			echo "$CAPTURE_BOOT_END $CAPTURE_ROOT_END" > $DISK_ROOT_MOUNT/origroot_boundaries
			
			umount $DISK_ROOT_MOUNT >> $IMAGE_LOG 2>&1
			sync
			
			resize_partition $TEMP_PART_SIZE_KB
			return $?

		#	If not, then exit out with a "Not enough space" error
		else
			echo "Not enough free disk space to perform imaging"
			RESULT=5
			exit 5
		fi
	fi
}

mountTempPartition() {
	[ $DEBUG -ne 0 ] && set -x
	# Create an ext4 FS and mount.
	mkfs.ext4 -m 0 -O ^has_journal -L "TEMP_PART" $TEMP_PART >>$IMAGE_LOG 2>&1
	mount $TEMP_PART $TEMP_MOUNT >> $IMAGE_LOG 2>&1
	RET=$?
	
	if [ $RET -ne 0 ]
	then
		# Disk was busy and the internal table in the kernel couldn't be updated. Force a reread.
		blockdev --rereadpt $FLASH_DEVICE
		mount $TEMP_PART $TEMP_MOUNT

		# If it failed again, now really exit
		if [ $? -ne 0 ]
		then
			cleanUp
			RESULT=5
			exit 5
		fi
	fi
}

prepareDDFile() {
	[ $DEBUG -ne 0 ] && set -x
	# Create a dd file for the new disk imge on the temporary partition
	dd if=/dev/zero of=$DD_FILE bs=1024 count=$DD_SIZE_KB >>$IMAGE_LOG 2>&1

	DISK_LOOP=$(losetup -f)
	losetup $DISK_LOOP $DD_FILE

	# copy first MB, should be enough for a while
	dd if=$FLASH_DEVICE of=$DISK_LOOP count=2048 >/dev/null 2>&1
	
	partition_dd $DISK_LOOP

	# Copy the mbr from the flash disk to a temp file
	dd if=$FLASH_DEVICE of=/tmp/mbr.dd bs=446 count=1 >>$IMAGE_LOG 2>&1
	sync
		
	# Setup Boot partition loopback
	BOOT_LOOP=$(losetup -f)
	losetup $BOOT_LOOP $DD_FILE -o $(( $(fdisk -l $DISK_LOOP 2>/dev/null| grep loop0p1 | sed -e 's/*//' | awk '{print $2}') * 512 ))

	# Setup Boot partition loopback
	ROOT_LOOP=$(losetup -f)
	losetup $ROOT_LOOP $DD_FILE -o $(( $(fdisk -l $DISK_LOOP 2>/dev/null| grep loop0p2 | sed -e 's/*//' | awk '{print $2}') * 512 ))
		
	# Copy boot partition from disk
	dd if=$BOOT_PART of=$BOOT_LOOP >>$IMAGE_LOG 2>&1
		
	mkfs.btrfs ${MKFS_BTRFS_OPTIONS} -d single -m single -L "ROOT" $ROOT_LOOP  >>$IMAGE_LOG 2>&1
	
	# Mount btrfs partitions... orig disk and new image
	[ -z "`mount | grep $DISK_ROOT_MOUNT`" ] && mount -o $BTRFS_MOUNT_OPTS $ROOT_PART $DISK_ROOT_MOUNT
	mount -o loop,rw,$BTRFS_MOUNT_OPTS $ROOT_LOOP $IMAGE_ROOT_MOUNT

	# Before removing loopbacks, insert the mbr.  Had to do this down here, something I was doing would wipe out my mbr.
	dd if=/tmp/mbr.dd of=$DISK_LOOP bs=446 count=1 >>$IMAGE_LOG 2>&1
	sync
}

copyImageData() {
	[ $DEBUG -ne 0 ] && set -x
	# Copy the data from the flash disk to our new dd image
	# Using root and writable as a base, we will prime our sending/receiving  loop by first sending over the permanent subvolumes 
	
	echo "Copying base subvolumes. This may take a few minutes..."
	for i in `btrfs subvol list $DISK_ROOT_MOUNT | grep -v "snapshots" | cut -d ' ' -f9` ; do 
		SNAP_NAME="$DISK_ROOT_MOUNT/${i}_ro"
	
		btrfs subvol snap -r $DISK_ROOT_MOUNT/${i} $SNAP_NAME >>$IMAGE_LOG 2>&1
		btrfs fi sync $DISK_ROOT_MOUNT >>$IMAGE_LOG 2>&1
		
		btrfs send $SNAP_NAME | btrfs receive $IMAGE_ROOT_MOUNT >>$IMAGE_LOG 2>&1
		btrfs fi sync $IMAGE_ROOT_MOUNT >>$IMAGE_LOG 2>&1
	done;
		
	# With root and writable in place, we can perform incremental sends using them as a base.
	echo "Copying snapshots. This may take a few minutes..."
	btrfs subvol list $DISK_ROOT_MOUNT | grep "snapshots" | cut -d ' ' -f9 | xargs -I {} bash -c '
	NAME="$DISK_ROOT_MOUNT/{}_ro_`echo {} | cut -d '/' -f2`"
	SUBVOL="`echo {} | cut -d '/' -f3`"
		
	# Create the snapshot
	btrfs subvol snap -r $DISK_ROOT_MOUNT/{} $NAME
	btrfs fi sync $DISK_ROOT_MOUNT
		
	# Send the snapshot
	btrfs send -p $DISK_ROOT_MOUNT/${SUBVOL}_ro $NAME | btrfs receive $IMAGE_ROOT_MOUNT
	btrfs fi sync $IMAGE_ROOT_MOUNT
	' >>$IMAGE_LOG 2>&1
		
	# With all the snapshots over, they now need to be recreated as read write in their proper locations
	echo "Cleaning up..."
	btrfs subvol list $IMAGE_ROOT_MOUNT | cut -d ' ' -f9 | xargs -I {} bash  -c '
	SNAP="`echo {} | cut -d '_' -f3-`"
	DEST="$IMAGE_ROOT_MOUNT"
		
	if [ -n "$SNAP" ]
	then
		DEST="${DEST}/snapshots/${SNAP}"
		mkdir -p $DEST
	fi
		
	# Create the snapshot
	btrfs subvol snap ${IMAGE_ROOT_MOUNT}/{} $DEST/`echo {} | cut -d '_' -f1`
	btrfs fi sync $IMAGE_ROOT_MOUNT
		
	# Delete the read only copy
	btrfs subvol del ${IMAGE_ROOT_MOUNT}/{}
	btrfs fi sync $IMAGE_ROOT_MOUNT
	' >>$IMAGE_LOG 2>&1

	# unmount the image file, so that its buffer are flushed
	# this prevents the file from being modified while we transfer it
	umount $IMAGE_ROOT_MOUNT
}

detectTableType(){
	[ $DEBUG -ne 0 ] && set -x
	fdisk -l $FLASH_DEVICE 2>/dev/null | tail -1 |grep -q GPT
	if [ $? -ne 0 ]; then
		echo "MBR"
	else 	
		echo "GPT"
	fi
}

capturePreCheck(){
	if [ -b ${FLASH_DEVICE}3 ]
	then
		# The extra partition exists. First, delete the partition
		parted -s $FLASH_DEVICE rm 3 >> $IMAGE_LOG
		
		# We put way the original partition boundaries on the root disk. Grab that, so we can restore the partition
		RSTR_DIR=$(mktemp -d /tmp/RSTR.XXXXX)
		
		mount -o $BTRFS_MOUNT_OPTS ${FLASH_DEVICE}2 $RSTR_DIR >> $IMAGE_LOG
		ORIG_BOUNDARIES=$(cat $RSTR_DIR/origroot_boundaries)
		
		umount $RSTR_DIR
		
		# Restore the partition boundaries
		parted -s $FLASH_DEVICE rm 2 >> $IMAGE_LOG
		if [ "x$tableType" = "xMBR" ] ; then
			MKPARTOPT="primary btrfs"
		else 
			MKPARTOPT="ROOT btrfs"
		fi
		parted -s $FLASH_DEVICE mkpart $MKPARTOPT "$ORIG_BOUNDARIES" >> $IMAGE_LOG 2>&1
		
		# Remount the partition to do a btrfs resize
		mount -o $BTRFS_MOUNT_OPTS ${FLASH_DEVICE}2 $RSTR_DIR >> $IMAGE_LOG
		btrfs fi resize 1:max $RSTR_DIR >> $IMAGE_LOG
		btrfs fi sync $RSTR_DIR >> $IMAGE_LOG
		
		# Delete the boundaries file
		rm $RSTR_DIR/origroot_boundaries
		
		umount -l $RSTR_DIR
		
		rmdir $RSTR_DIR
	fi
}

capture_image(){
	[ $DEBUG -ne 0 ] && set -x
	# Check If transferImageFile has been defined
	type -t transferImageFile > /dev/null 2>&1
	if [ $? -ne 0 ]
	then
		echo "The function \"transferImageFile\" has not been implemented!"
		echo "Please implement this function and then try again."
		RESULT=9
		exit 9
	fi
	
	# Capture Precheck. If during a capture, it abnormally terminates (such as due to a power failure), the scratch partition may be left behind.
	# Check for its existance here before we start doing anything and delete it if it exists. Restore the root partition back as well.

	tableType=`detectTableType`
	
	capturePreCheck
	
	echo "Checking Free Disk Space"
	checkFreeDiskSpace
	if [ $? -ne 0 ] ; then
	        echo "Not enough space"
	        cleanUp
		RESULT=5
	        exit 5
	fi

	echo "Creating Temporary Partition"
	TEMP_PART=$(create_temp_partition $FLASH_DEVICE)
	
	echo "Mounting Temporary Partition"
	mountTempPartition
	
	echo "Preparing DD File"
	prepareDDFile

	echo "Copying Image Data"
	copyImageData

	type -t customTransferImageFile > /dev/null 2>&1
	if [ $? -ne 0 ]
	then
		echo "Upload image using default transferImageFile()"
		transferImageFile
	else
		echo "Upload image using customTransferImageFile()"
		customTransferImageFile $DD_FILE
		[ $? -ne 0 ] && RESULT=8
	fi

	cleanUp

	return $RESULT
}
