#!/bin/sh

[ ! -e /dev/$1 ] && echo "Must specify device" && exit

# Get the size of the s1 slice of that device, in MB, plus two
# (gpart outputs in 512 byte blocks) 
expr `gpart show|grep ${1}s1|awk '{print $3}'` "/" 2048 "+" 2 > tmp
megs=`cat tmp`
echo "/dev/${1}s1 size: $megs MB"

if [ -e clean -a image -nt install ]; then
	echo "install script not changed, copying existing image to /dev/$1"
	dd if=image of=/dev/$1 bs=1M
	exit
fi

if [ ! -e clean ]; then
	echo "File 'clean' not found. Assuming there is a fresh FreeBSD"
	echo "installation on device /dev/$1"
	echo "Press enter to copy this to file 'clean' in this directory"
	read pause
	dd if=/dev/$1 of=clean bs=1M count=$megs
	[ -e image ] && rm image
else
	echo "copying clean image to /dev/$1"
	dd if=clean of=/dev/$1 bs=1M
fi

# Label the slice and set it up so it will boot no matter what the 
# disk device is called in the system that has to boot from it later.
glabel label boot /dev/${1}s1
mount /dev/label/boota /mnt
echo "/dev/label/boota / ufs rw 1 1" > /mnt/etc/fstab

# Turn off boot menu
echo 'autoboot_delay="-1"' >> /mnt/boot/loader.conf
echo 'beastie_disable="YES"' >> /mnt/boot/loader.conf

# Turn off syslogd and devd while still setting up
# (no need for them, and pipes in /var/run won't copy)
cp /mnt/etc/rc.conf /mnt/etc/rc.conf.orig
echo 'syslogd_enable="NO"' >>/mnt/etc/rc.conf
echo 'devd_enable="NO"' >>/mnt/etc/rc.conf

# Backup /etc/gettytab and create new terminal-type that auto-logins as root
cp /mnt/etc/gettytab /mnt/etc/gettytab.orig
echo "Rl|rootlogin console:ht:np:sp#115200:al=root" >> /mnt/etc/gettytab

# Back up /etc/ttys, remove ttyv0 entry and stick in 
# a new line that uses our auto-login entry
mv /mnt/etc/ttys /mnt/etc/ttys.orig
grep -v ttyv0 /mnt/etc/ttys.orig > /mnt/etc/ttys
echo 'ttyv0	"/usr/libexec/getty Rl" cons25 on secure' >> /mnt/etc/ttys

# Write to .login so that installation starts when root logs in
mv /mnt/root/.login /mnt/root/.login.orig
echo "/root/FreeWDE-install" > /mnt/root/.login
touch /mnt/root/.hushlogin


# The remainder of this script creates a file,
# named 'FreeWDE-install' in /root.


########## FreeWDE-install ##########

cat > /mnt/root/FreeWDE-install <<-'EOF'
	#!/bin/sh

	# getsize <variable> <minsize> <allow_all> <prompt>
	# Get a partition size in G or M, see if it is at least <minsize> MB
	# Allow input 'all' to output "*" if <allow_all> is 'true'
	# Remainder is used as a prompt

	lastchar() { echo $1|awk '{print substr($0, length($0), 1)}'; }
	numberchars() { echo $1|awk '{gsub(/[^0-9]/, ""); print $0}'; }
	getsize()
	{
		unset out
	    variable=$1
	    minsize=$2
	    all=$3
	    shift 3
	    prompt=$*
	    while [ "xx$out" = xx ]; do
	        read -p "$prompt: " in
	        if [ "xx$in" != xx ]; then
	            if [ "$in" = "all" -a $all = "true" ]; then
	                out=*
	            else
	                if [ "$in" = 0 -o "$in" = "none" ]; then
	                	in="0m"
	                fi
	                lc=`lastchar "$in"`
	                in=`numberchars $in`
	                if [ $lc = "G" -o $lc = 'g' ]; then
	                    mb=`expr $in "*" 1024`
	                    out=${in}G
	                elif [ $lc = "M" -o $lc = 'm' ]; then
	                    mb=$in
	                    out=${in}M
	                    [ $out = 0M ] && out=0
	                fi
	                if [ -z $out ]; then
	                    echo "Must be a number followed by 'M' or 'G'"
	                elif [ $mb -lt $minsize ]; then
	                    unset out
	                    echo "You must enter at least $minsize megabytes."
	                fi
	            fi
	        fi
	    done
	    setvar $variable "$out"
	}
	
	# Returns 0 or 1 depending on whether user enters 'y'/'yes' or 'n'/'no'
	yesno()
	{
	    while [ 1 ]; do
	        read -p "$* (yes/no) " in
	        if [ "xx$in" != xx ]; then
	            [ "$in" = "yes" -o "$in" = "y" ] && return 0
	            [ "$in" = "no" -o "$in" = "n" ] && return 1
	        fi	
	    done
	}
	
	
	# Find the disk device underneath the label we booted from
	dev=/dev/`glabel list | sed -n '/boot/,/Name:/ p' | sed -n 's/1. Name: \(.*\)s1/\1/p'`
	if [ ! -e $dev ]; then
	    echo "Error: Cannot find boot device"
	    exit
	fi
	
	conscontrol mute on >/dev/null
	clear
	
	if [ ! -e phase2 ]; then
		
		# Talk to user
		cat <<-'END'
			Welcome to FreeWDE, the FreeBSD with Whole Disk Encryption installer
			
			This script will help you create a bootable disk with multiple "slices". One
			will be unencrypted and will hold the files needed to boot. In fact, we will
			reuse the part that you just booted from and just keep it as it is.
			
			Then there's the encrypted slice that this is all about. It is encrypted with
			AES-256 and holds either one or two partitions. There's at least the root
			filesystem, and if you so choose there is also a swap area. And to make for a
			faster system on slow disks (such as many USB sticks), /tmp and /var/log can
			be put on tmpfs (which means they usually stay in RAM).
			
			You can choose to install another 'unencrypted' slice. In there, we'll put
			a regular Windows FAT32 filesystem. This means you can use the device as a 
			regular USB disk or SD-card. To stick in your camera and take pictures on, for
			example.
	
		END
		
		yesno "Are you sure you want to do this?" || exit
			
		cat <<-'END'

	
			Enter size of UNENCRYPTED slice. This size can be entered in megabytes or
			gigabytes. It needs to be entered as a number followed by M or G. So '1G'
			for a 1 gigabyte unencrypted slice. Enter "0" or "none" if you do not want
			an unencrypted slice.
				
		END
			
		getsize unencryptedsize 0 false "Size of unencrypted slice"
		
		# existence of file 'unencrypted' is signal for phase2 to format it
		if [ $unencryptedsize != 0 ]; then
			echo "p 2 11 * ${unencryptedsize}" >> fdisk-input
			touch unencrypted
		fi
		
		
		cat <<-'END'
		
		
			Enter the size of the ENCRYPTED slice: a number immediately followed
			by M or G. This size includes root file system and swap.
			If you enter "all", we will use the remainder of the device.
			
		END
		
		getsize size 500 true "Size of encrypted slice"
	
		echo "p 3 165 * ${size}" >> fdisk-input
	
		cat <<-'END'
	
	
			How big is the swap partition on the encrypted slice, again as a number
			immediately followed by the capital letter M or G.  If you enter "0" you
			will not have swap, so the entire encrypted slice is used for files.
			
		END
	
		getsize swapsize 0 false "Size of swap space"
		
		# file 'swap' contains size, for phase2
		[ $swapsize != 0 ] && echo "$swapsize" > swap
		
		cat <<-'END'
		
		
			You can choose to wipe the disk, putting zeroes in the unencrypted slice and
			random data in the encrypted one. This makes sure there is no previous data
			that can be read, but it can take a long time depending on size and speed
			of the disk. So if you just bought it, you might as well enter no here.
		
		END
		
		yesno "Wipe disk?" && touch wipe

		echo ""
		echo ""
		yesno "Would you like /tmp on a tmpfs (in ram) ?" && touch tmpram		
		
		echo ""
		echo ""
		echo "We can do the same for /var/log. This does mean you'll have to edit /etc/fstab"
		echo "before you can debug why something crashes."
		echo ""
		yesno "Would you like /var/log on a tmpfs?" && touch logram
		
		echo ""
		echo ""
		echo "writing partition table"
		
		# Allow us to write to partition table of current disk
		sysctl kern.geom.debugflags=16 2>&1 >/dev/null
		
		echo ""
		echo ""
		echo "The partition table needs to be re-read, for which we need to reboot."
		echo "Press enter to reboot. Installation will continue after we return."
		read pause

		# Write partition table
		fdisk -f fdisk-input $dev 2>&1 >/dev/null
		clear
		
		touch phase2
		
		reboot
	
	else

		echo "FreeWDE - FreeBSD with Whole Disk Encryption"
		echo "Continuing the installation..."
	
		# Create second labeled device
		glabel label crypt ${dev}s3 >/dev/null
		cryptdev=/dev/label/crypt
		
		if [ -e wipe ]; then
			# Fill slice with random
			echo ""
			echo "Now writing random data to the new crypto slice. This might take a while."
			echo "(Pressing Ctrl-T will tell you how many bytes have been written.)"
			echo ""
			dd if=/dev/random of=$cryptdev bs=1M 2>&1 >/dev/null
		fi
		
		# Create crypto disk
		echo ""
		echo "This is where you pick a really strong passphrase and enter it twice"
		echo "to create the encrypted disk."
		echo ""
		geli init -B none -b -e aes -l 256 -s 512 $cryptdev
	
		# Attach to the crypto disk
		echo ""
		echo "Now type your passphrase once more to attach to the new encrypted disk."
		echo ""
		while [ ! -e /dev/label/crypt.eli ]; do
			geli attach $cryptdev
		done
		
		# Make partition(s) on this slice.
		echo ""
		echo "Partitioning"
		bsdlabel -w ${cryptdev}.eli  2>&1 >/dev/null
		echo "8 partitions:" > bsdlabel-input
		echo "  a: * * 4.2BSD 4096 16384" >> bsdlabel-input
		[ -e swap ] && echo "  b: `cat swap` * swap" >> bsdlabel-input
		
		bsdlabel $cryptdev.eli | grep "c:" >> bsdlabel-input
		bsdlabel -R ${cryptdev}.eli bsdlabel-input 2>&1 >/dev/null
		
		# format the partition and mount it as /mnt/crypt
		echo ""
		echo "Formatting"
		newfs ${cryptdev}.elia 2>&1 >/dev/null
		mkdir /mnt/crypt
		mount ${cryptdev}.elia /mnt/crypt 2>&1 >/dev/null || exit 1
		
		# Copy files to encrypted disk
		# Skipping:
		#    /mnt because of the recursion issue
		#    /rescue because it's full of hardlinks
		#    /boot because we're symlinking that to the unencrypted part
		echo ""
		echo "Copying files:"
		for i in `ls /`; do
			if [ $i != mnt -a $i != rescue -a $i != boot ]; then
				echo "  /$i"
				if [ $i != usr ]; then
					cp -Rp /$i /mnt/crypt 2>&1 >/dev/null
				else
					# Special treatment for /usr because it takes very long
					# So we show separate subdirs. Go ahead, call me vain..
					mkdir /mnt/crypt/usr
					for j in `ls /usr`; do
						echo "    /$j"
						cp -Rp /usr/$j /mnt/crypt/usr 2>&1 >/dev/null
					done
				fi
			fi
		done
		
		# link to /boot in unencrypted part as /boot in encrypted system
		echo ""
		echo "Creating mount points and symlinks"
		mkdir /mnt/crypt/mnt
		mkdir /mnt/crypt/mnt/boot
		ln -s /mnt/boot/boot /mnt/crypt/boot
		[ -e unencrypted ] && mkdir /mnt/crypt/unencrypted
		
		# Set things up to switch to the encrypted disk next time it boots
		echo ""
		echo "Making /etc/fstab"
		mv /etc/fstab /etc/fstab.orig
		[ -e swap ] && echo "${cryptdev}.elib none swap sw 0 0" >> /etc/fstab
		echo "${cryptdev}.elia / ufs rw 1 1" >> /etc/fstab
		echo "/dev/label/boota /mnt/boot ufs rw 1 1" >> /etc/fstab
		[ -e unencrypted ] && echo "${dev}s2 /unencrypted msdosfs rw 0 2" >> /etc/fstab
		[ -e tmpram ] && echo "tmpfs /tmp tmpfs rw 0 0" >> /etc/fstab
		[ -e logram ] && echo "tmpfs /var/log tmpfs rw 0 0" >> /etc/fstab
		
		# format FAT32 slice if we created one
		if [ -e unencrypted ]; then
			
			if [ -e wipe ];then
				echo ""
				echo "Wiping the unencrypted slice"
				dd if=/dev/zero of=${dev}s2 bs=1M  2>&1 >/dev/null
			fi
	
			echo ""
			echo "Formatting unencrypted slice"
			newfs_msdos -F32 ${dev}s2 2>&1 >/dev/null
		fi
		
		cp /etc/fstab /mnt/crypt/etc/fstab
		
		echo 'geom_eli_load="YES"' >> /boot/loader.conf
		
		# Take out the auto-login for root
		mv /mnt/crypt/etc/ttys.orig /mnt/crypt/etc/ttys
		mv /mnt/crypt/etc/gettytab.orig /mnt/crypt/etc/gettytab
		
		# Put back the empty rc.conf
		mv /mnt/crypt/etc/rc.conf.orig /mnt/crypt/etc/rc.conf
				
		# take stuff out of /root
		mv /mnt/crypt/root/.login.orig /mnt/crypt/root/.login
		rm /mnt/crypt/root/.hushlogin
		
		rm /mnt/crypt/root/*
	
		echo ""
		echo "Press enter to boot into encrypted system"
		read pause
		
		# reboot into encrypted system
		echo ""
		echo "Rebooting ..."
		echo ""
		reboot

	fi
	
EOF

chmod +x /mnt/root/FreeWDE-install
umount /mnt

echo "copying modified image back"
dd if=/dev/$1 of=image bs=1M count=$megs