#!/bin/sh
#  Copyright, 2012 Dustin Kirkland <kirkland@ubuntu.com>
#  Copyright, 2012 Scott Moser <smoser@ubuntu.com>
#  Copyright, 2012 Axel Heider
#
#  Based on scripts from
#    Sebastian P.
#    Nicholas A. Schembri State College PA USA
#    Axel Heider
#    Dustin Kirkland
#    Scott Moser
#
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see
#    <http://www.gnu.org/licenses/>.

# generic settings
# ${ROOT} and ${rootmnt} are predefined by caller of this script. Note that
# the root fs ${rootmnt} it mounted readonly on the initrams, which fits
# nicely for our purposes.
root_rw=/media/root-rw
root_ro=/media/root-ro
ROOTMNT=${rootmnt:-/root} # use global name to indicate created outside this
logpath=${root_rw}/overlay.log
backuppath=${root_rw}/backup
dir_prefix="overlay"
upperdir=${root_rw}/${dir_prefix}
flagpath=${upperdir}/home/.config

prereqs=""
case "$1" in
	# no pre-reqs
	prereqs) echo ""; exit 0;;
esac

if [ "$(cat /proc/cmdline | grep overlay)" = "" ]; then
	if [ -f "${ROOTMNT}/etc/default/grub" ]; then
            if [ "$(cat ${ROOTMNT}/etc/default/grub | grep overlay)" ]; then
                sed -i "s/overlay //g" ${ROOTMNT}/etc/default/grub
                if [ $? -ne 0 ];then
                    echo "fail"
                    exit 0
                fi
            fi
	fi
	exit 0
fi

# check if kernel module exists
insmod ${ROOTMNT}/usr/lib/modules/$(uname -r)/kernel/fs/overlayfs/overlay.ko*

. /scripts/functions
. "${ROOTMNT}/etc/kylinoverlay.conf"
. "${ROOTMNT}/etc/kylinoverlayuser.conf"

#if [ "$kylin_debug" = "1" ]; then
#	echo "----------$(date)----------" >> $logpath
#else
#	echo "----------$(date)----------" 
#fi

log() {
	if [ "$kylin_debug" = "1" ]; then
		echo "$1" >> $logpath
	else
		echo "$1"
	fi
}
#log "ROOTMNT = $ROOTMNT"

overlayrootify_fstab() {
	# overlayrootify_fstab(input, root_ro, root_rw, dir_prefix, recurse, swap)
	# read input fstab file, write an kylinoverlay version to stdout
	# also returns (_RET) a list of directories that will need to be made
	local input="$1" root_ro="${2:-/media/root-ro}"
	local root_rw="${3:-/media/root-rw}" 
	local recurse=${4:-1} swap=${5:-0}
	local hash="#" oline="" ospec="" upper="" dirs="" copy_opts=""
	local vfstypes=" proc sys tmpfs dev udev "
	local spec file vfstype opts pass freq line ro_line

	[ -f $input ] || return 1
	while read spec file vfstype opts pass freq; do
		line="$spec $file $vfstype $opts $pass $freq"
		case ",$opts," in
			*,ro,*) ro_opts="$opts";;
			*) ro_opts="ro,${opts}";;
		esac
		ro_line="$spec ${root_ro}$file $vfstype ${ro_opts} $pass $freq"
		if [ "$file" = "/data" ]; then
			echo "#${line}"
		else
			echo "$line"
		fi
	done < "$input"
}

mksnapshot() {
	if [ ! -d ${backuppath} ]; then
		mkdir -p "${backuppath}"
	fi

	if [ -f ${flagpath}/$1 ]; then
		initstring=$1
		suffix=${initstring##*_}
		name=$2
		log "rm -rf ${flagpath}/$1"
		rm -rf "${flagpath}/$1"
	else
		log "miss ${flagpath}/$1"
	fi

	if [ -f ${backuppath}/index ];then
		updatenum=$(cat "${backuppath}/index")
	else
		log "miss ${backuppath}/index"
	fi

	if [ -d ${upperdir} ];then
		updatenum=$((updatenum+1))
		if [ $updatenum -gt $max_update_num ]; then
			updatenum=$((updatenum - 1))
			exit 0
		fi
		if [ "${backuppath}/update_${updatenum}*" ]; then
			rm -rf "${backuppath}/update_${updatenum}*"
		fi
		log "mv "${upperdir}" "${backuppath}/update_${updatenum}_${suffix}_${name}""
		mv "${upperdir}" "${backuppath}/update_${updatenum}_${suffix}_${name}" 
		mkdir -p "${upperdir}"
		echo $updatenum > "${backuppath}/index"
	fi
	####
}

backup_stash() {
		mkdir -p "${BACKUP_PATH}"
		log "BACKUP_TARGET ${BACKUP_TARGET}"
		for backup_target in ${BACKUP_TARGET}
		do
			log "backup $backup_target"
			if [ -e ${upperdir}${backup_target} ]; then
				backup_dir=$(dirname ${upperdir}${backup_target})
					log "mkdir -p ${BACKUP_PATH}${backup_dir}"
					mkdir -p ${BACKUP_PATH}${backup_dir} >> $logpath 2>&1
					log "mv ${upperdir}${backup_target} ${BACKUP_PATH}${upperdir}${backup_target}"
					rsync -a ${upperdir}${backup_target} ${BACKUP_PATH}${backup_dir}  >> $logpath 2>&1
			fi
		done
}
 
backup_pop() {
		log "backup_pop"
		for backup_target in ${BACKUP_TARGET}
		do
			backup_dir=$(dirname ${upperdir}${backup_target})
			log "mkdir -p ${backup_dir}"
			mkdir -p ${backup_dir}  >> $logpath 2>&1
			log "mv ${BACKUP_PATH}${upperdir}${backup_target} ${upperdir}${backup_target}"
			rsync -a ${BACKUP_PATH}${upperdir}${backup_target} ${backup_dir} >> $logpath 2>&1
		done
		#rm -rf ${BACKUP_PATH}
		rm -rf ${flagpath}/mk_rollback*
}

rollback() {
	backup_stash
	case "$1" in
		last)
			##
			if [ -d ${upperdir} ]; then
				log "rm -rf "${upperdir}""
				rm -rf "${upperdir}"
			fi
			##
			if [ -f ${backuppath}/index ];then
				updatenum=$(cat "${backuppath}/index")
				
				if [ -d ${backuppath}/update_${updatenum}* ]; then
					log "mv "${backuppath}/update_${updatenum}*" "${upperdir}""
					mv "${backuppath}/update_${updatenum}"* "${upperdir}" >> $logpath
				else
					log "miss ${backuppath}/update_${updatenum}"
				fi
				
				updatenum=$((updatenum-1))
				if [ -d ${backuppath}/update_${updatenum}* ];then
					echo $updatenum > "${backuppath}/index"
				else
					log "rm -rf "${backuppath}/index""
					rm -rf "${backuppath}/index"
				fi
			fi
			;;
		*)
			log "rollback action error!"
			;;
	esac
	backup_pop
}

getlowdir() {
	log "GETLOWDIR"
	if [ -d ${backuppath}/update_1* ]; then
		for update_dir_name in ${backuppath}/update_*
		do
			lowdir_name=${update_dir_name}:${lowdir_name}
		done
	else
		lowdir_name=""
	fi
	lowdir_name=${lowdir_name}${root_ro}
}

set_mount_label() {
	local input="$1" mount_point
        case "$input" in
			LABEL=*)
				mount_point="/dev/disk/by-label/${input#LABEL=}"
				;;
			PARTLABEL=*)
				mount_point="/dev/disk/by-label/${input#PARTLABEL=}"
				;;
			UUID=*) 
				mount_point="/dev/disk/by-uuid/${input#UUID=}" 
				;;
			/dev/*) 
				mount_point="${input}"
				;;
            *) 
				mount_point="/dev/${input}"
				;;
        esac
	echo "${mount_point}"
}

read_fstab() {
	local input=$1
	[ -f $input ] || return 1
	while read spec file vfstype opts pass freq; do
		if [ "$file" = "/" ]; then
			kylinoverlay=$(set_mount_label $spec)
		elif [  "$file" = "/data" ]; then
			kylinupoverlay=$(set_mount_label $spec)
		elif [ "$file" = "/backup" ]; then
			kylindownload=$(set_mount_label $spec)
		elif [ "$file" = "/boot" ]; then
			kylinboot=$(set_mount_label $spec)
		elif [ "$file" = "/boot/efi" ]; then
			kylinefi=$(set_mount_label $spec)
		fi
	done < "$input"
}


read_fstab "${ROOTMNT}/etc/fstab"
# make the mount point on the init root fs ${root_rw}
mkdir -p "${root_rw}" 
# make the mount point on the init root fs ${root_ro}
mkdir -p "${root_ro}"

/sbin/mount ${kylinupoverlay} ${root_rw} >> $logpath 2>&1

log "--------------------" 

mkdir -p "${upperdir}" ||
	log "failed to create ${upperdir}"

updatenum=0

if [ -d ${flagpath} ]; then
	for file in ${flagpath}/*; do
		tmp=${file##*/}
		suf=${tmp%_*}
		if [ "${tmp}" != "." -a "${tmp}" != ".." ]; then
			name=$(cat ${flagpath}/${tmp})
			if [ "$suf" = "mk_snapshot" ]; then
				mksnapshot $tmp $name
			fi
			if [ "$suf" = "mk_rollback" ]; then
				action=`cat ${flagpath}/${tmp}`
				rollback $action
			fi
		fi
	done
else
	log "miss ${flagpath}"
fi
if [ -d ${backuppath} ]; then
	for file in ${backuppath}/*; do
 		tmp=${file##*/}
        	if [ "$tmp" = "rollback_deepfreeze" ];then
        		rm -rf ${upperdir}/*
        	fi
	done
else
	log "miss ${backuppath}"
fi

getlowdir

# settings based on kylinoverlay_driver
mount_type="overlay"
mount_opts="-o index=off -o metacopy=off -o lowerdir=${lowdir_name},upperdir=${upperdir}"
# 3.18+ require 'workdir=' option.
case "$(uname -r)" in
	2*|3.1[01234567]*|3.[0-9].*) :
			;;
	*)
		workdir="${dir_prefix%/}-workdir"
		mount_opts="${mount_opts},workdir=${root_rw}/${workdir}"
		mkdir -p ${root_rw}/${workdir}
		;;
esac
mount_opts="${mount_opts} $kylinoverlay ${ROOTMNT}"

log "mount --move ${ROOTMNT} ${root_ro}"
log "mount -t $mount_type $mount_opts"
# root is mounted on ${ROOTMNT}, move it to ${ROOT_RO}.
# there is nothing left at ${ROOTMNT} now. So for any error we get we should
# either do recovery to restore ${ROOTMNT} for drop to a initramfs shell using
# "panic". Otherwise the boot process is very likely to fail with even more 
# errors and leave the system in a wired state. 
/sbin/mount --move ${ROOTMNT} ${root_ro} 2>&1
if [ $? -ne 0 ]; then
	log "fail to mount move rootmnt to root_ro!"
	exit 0
else
	/sbin/mount -t $mount_type $mount_opts >> $logpath 2>&1
	if [ $? -ne 0 ]; then
       		log "failed to create new ro/rw layered ${ROOTMNT}"
       		log "mount command was: mount -t $mount_type $mount_opts"
       		# do recovery and try resoring the mount for ${ROOTMNT}
       		/sbin/mount --move ${root_ro} ${ROOTMNT}
      		if [ $? -ne 0 ]; then
              		# thats bad, drop to shell to let the user try fixing this
               		log "RECOVERY_ERROR: failed to move $root_ro back to ${ROOTMNT}"
       		fi
        	exit 0
	fi
fi

#/sbin/mount ${kylinboot} "${ROOTMNT}/boot"
#/sbin/mount ${kylinefi} "${ROOTMNT}/boot/efi"

# technically, everything is set up nicely now. Since ${ROOTMNT} had beend
# mounted read-only on the initfamfs already, ${ROOTMNT}${root_ro} is it,
# too.  Now we init process could run - but unfortunately, we may have to
# prepare some more things here. 
# Basically, there are two ways to deal with the read-only root fs. If the
# system is made aware of this, things can be simplified a lot.  If it is
# not, things need to be done to our best knowledge. 
#
# So we assume here, the system does not really know about our read-only
# root fs.
#
# Let's deal with /etc/fstab first. It usually contains an entry for the
# root fs, which is no longer valid now. We have to remove it and add our
# new ${root_ro} entry. 
# Remember we are still on the initramfs root fs here, so we have to work
# on ${ROOTMNT}/etc/fstab. The original fstab is
# ${ROOTMNT}${root_ro}/etc/fstab.
cat <<EOF >${ROOTMNT}/etc/fstab
#
#  This fstab is in an overlayfs.  The real one can be found at
#  ${root_ro}/etc/fstab
#  The original entry for '/' and other mounts have been updated to be placed
#  under $root_ro.
#  To permanently modify this (or any other file), you should change-root into
#  a writable view of the underlying filesystem using:
#      sudo kylin-chroot
#
EOF
[ $? -eq 0 ] || log "failed to modify /etc/fstab (step 1)"
overlayrootify_fstab ""${root_ro}/etc/fstab"" "$root_ro" \
	"$root_rw"  "$recurse" "$swap" >> "${ROOTMNT}/etc/fstab" ||
	log "failed to modify /etc/fstab (step 2)"
	umount ${root_rw} 2>&1
	umount ${root_ro} 2>&1
exit 0
