#!/bin/bash
# This script allows backuping a complete SD card, NAND or eMMC card using tar utility
#

_resolve_file_location(){
	SOURCE="$1"
	while [ -h "$SOURCE" ]; do # resolve $SOURCE until the file is no longer a symlink
	  THE_DIR="$( cd -P "$( dirname "$SOURCE" )" && pwd )"
	  SOURCE="$(readlink "$SOURCE")"
	  [[ "$SOURCE" != "/*" ]] && SOURCE="$DIR/$SOURCE" # if $SOURCE was a relative symlink, we need to resolve it relative to the path where the symlink file was located
	done
	THE_DIR="$( cd -P "$( dirname "$SOURCE" )" && pwd )"
	echo $THE_DIR
}
_get_relative_path_to(){
	source=$1
	target=$2
	common_part=$source
	back=
	while [ "${target#$common_part}" = "${target}" ]; do
	  common_part=$(dirname $common_part)
	  back="../${back}"
	done
	
	echo ${back}${target#$common_part/}
}
printHelp(){
	cat << EOF
NAME:
	diskimg -- Creates or restores a partition

SYNOPSIS:
	$0 -c TARGET_DIR -d DISK_DEVICE [-v]
	$0 -s SOURCE_DIR -d DISK_DEVICE [-v]
	
DESCRIPTION:
	Creates a backup of the given device using the `-c` option. All partition 
	table, boot sector, partition contents are stored into the given target
	directory in different files.
	
	Restores the backup using the `-r` option to the given device. All partition 
	table, boot sector, partition contents are restored from the given target
	directory
	

OPTIONS:
	-c	creates the backup using the given TARGET_DIR value
	-d	the DISK_DEVICE from/to process
	-e	exclude file patterns from tar (creation). Re-use for several values (e.g -e /foo/bar -e /my/path)
	-u	uncompressed image (disables compression)
	-z	uses gzip compression
	-j	uses bzip2 compression
	-P	keep existing partitions
	-r	restores from the data stored into the SOURCE_DIR value
	-J	uses XZ compression (DEFAULT)
	-T	skip tar archive creation (useful if backup only partition layout)
	-v	Verbose
	-v	Very verbose (debug mode)
	-D	dry run
	
EXAMPLES:
	One can easily recreate the partition scheme using using the following command 
	(which will not write the data from the tar files) :
		
		$0 -r /my/path/to/backup -d /dev/sdb -T
		
	Then it is possible to change the layout (resize filesystem to partition size e.g.)
	
	And then to restore the data (option -P will preserve existing partitions):
	
		$0 -r /my/path/to/backup -d /dev/sdb -P
	
EOF
exit 0
}

OPTS=$(getopt c:d:e:FDhjJPr:uTvVz $*)
#Test les paramètres
if [ $? != 0 ]
then
	echo "Error while retrieving paramger data from getopt"
	printHelp
    exit 1
fi

eval set -- "$OPTS"

DISK_DEVICE=""
TARGET_DIR=""
SOURCE_DIR=""
COMPRESSION_ALGORITHM="-J"
COMPRESSION_SUFFIX=".xz"
VERBOSE_OPTION=""
EXCLUDE_PATTERNS=""
KEEP_PARTITIONS=""
FAKE_DEVICE=""

while true ; do
    case "$1" in
        -d) DISK_DEVICE="$2" ; shift ; shift
        	;;
        -c) TARGET_DIR="$2" ; shift ; shift
        	;;
        -e) EXCLUDE_PATTERNS="$EXCLUDE_PATTERNS $2" ; shift ; shift
        	;;
        -F) FAKE_DEVICE="1"; shift
        	;;
        -r) SOURCE_DIR="$2" ; shift ; shift
        	;;
        -u) COMPRESSION_ALGORITHM="" ; COMPRESSION_SUFFIX="" ; shift
        	;;
        -z) COMPRESSION_ALGORITHM="-z" ; COMPRESSION_SUFFIX=".gz" ; shift
        	;;
        -j) COMPRESSION_ALGORITHM="-j"; COMPRESSION_SUFFIX=".bzip2" ; shift
        	;;
        -v) VERBOSE="1"; VERBOSE_OPTION="-v" ; shift
        	;;
        -V) VERBOSE="1"
        	VERBOSE_OPTION="-v"
        	set -x
        	shift
        	;;
        -D) DEBUG="1"; shift
        	;;
        -P) KEEP_PARTITIONS="1"; shift
        	;;
        -T) SKIP_TAR="1"; shift
        	;;
        -h|--help)
        	printHelp
        	exit 0
        	;;
        --) shift; break
        	;;
    esac
done

create_exclude_string(){
	EXCLUDE_OPTION="--exclude=."
	EXCLUDE_STRING=""
	for pattern in $(echo $EXCLUDE_PATTERNS) ; do 
		EXCLUDE_STRING="$EXCLUDE_STRING $EXCLUDE_OPTION$pattern"
	done
}

log_debug(){
	if [[ "$DEBUG" = "1" ]] ; then
		echo $*
	fi
}

prepare_btrfs_subvols(){
	local target_btrfs_dev=$1
	
	local temp_mount_point=$(mktemp -d)
	
	mount $target_btrfs_dev $temp_mount_point
	mkdir $temp_mount_point/.snapshots
	btrfs sub snap $temp_mount_point/ $temp_mount_point/.snapshots/ACTIVE
	#Gets the active sub snapshot id:
	local snap_id=$(btrfs sub list $temp_mount_point |grep '.snapshots/ACTIVE' | awk '{print $2}')
	
	#Sets it at the current snap
	btrfs sub set $snap_id $temp_mount_point
}

format_partition(){
	local partition_type="$1"
	local target_device="$2"
	echo "Formating $target_device with type $partition_type"
	
	if [[ "1" != "$DEBUG" ]] ; then
		case $partition_type in
		btrfs) mkfs.btrfs -f -L $partition_type $target_device
			test "$?" == 0 || exit $?
			prepare_btrfs_subvols $target_device
			;;
		*)
			mkfs -t $partition_type $target_device
			test "$?" == 0 || exit $?
			;;
		esac
	fi
	_REAL_TARGET_DEVICE=$target_device
}

create_luks_container(){
	local target_device="$1"
	
	echo "Creating a new LUKS container"
	if [ "$DEBUG" != "1" ] && [ "$FAKE_DEVICE" != "1" ] ; then
		cryptsetup luksFormat -c aes-xts-plain64 -s 512 -h sha512 $target_device <<EOF
$_luks_password
EOF
	fi
}

prepare_partition(){
	local partition_type="$1"
	local the_target_device="$2"
	local backup_basename=$3
	log_debug "partition_type=$partition_type"
	
	for part_type in $(echo $partition_type) ; do
		log_debug "part_type=$part_type"
		log_debug "the_target_device=$the_target_device"
		
		case $part_type in
		crypto_LUKS)
			ask_password $the_target_device confirm
			create_luks_container $the_target_device
			open_luks_container $the_target_device $backup_basename
			#Now that the partition is opened, we will format it at next loop:
			the_target_device=$_UNCRYPTED_DEVICE
			;;
		*)
			format_partition $part_type $the_target_device 
			;;
		esac
	done
	_REAL_FILESYSTEM_TYPE=$part_type
}

restore_partition(){
	local partition_file=$1
	local partition_dev=$2
	local backup_basename=$3
	
	mount_dir_name="diskimg-$backup_basename"
	mountpoint_dir="/mnt/$mount_dir_name"
	test "1" = "$VERBOSE" && echo "Mounting $partition_dev to $mountpoint_dir"
	if [[ "$FAKE_DEVICE" != "1" ]] ; then
		mkdir -p $mountpoint_dir
		test "$?" == 0 || exit $?
		mount $partition_dev $mountpoint_dir
		test "$?" == 0 || exit $?
	fi
	test "1" = "$VERBOSE" && echo "Restoring from $partition_file to $mountpoint_dir"
	
	if [[ "$SKIP_TAR" != 1 ]] ; then
		if [[ "1" = "$DEBUG" ]] ; then
			echo tar --numeric-owner $COMPRESSION_ALGORITHM $VERBOSE_OPTION -pxf $partition_file -C $mountpoint_dir
		fi
		if [[ "$FAKE_DEVICE" != "1" ]] ; then
			if [[ "1" != "$DEBUG" ]] ; then
				tar --numeric-owner $COMPRESSION_ALGORITHM $VERBOSE_OPTION -pxf $partition_file -C $mountpoint_dir
				test "$?" == 0 || exit $?
			fi
		fi
	fi
	
	test "1" = "$VERBOSE" && echo "Un-mounting $partition_dev"
	if [[ "$FAKE_DEVICE" != "1" ]] ; then 
		umount $partition_dev
		test "$?" == 0 || exit $?
		rmdir $mountpoint_dir
		test "$?" == 0 || exit $?
	fi
}

do_restore_disk(){
	echo "Restoring disk image from $SOURCE_DIR to $DISK_DEVICE"
	if [[ "$KEEP_PARTITIONS" != "1" ]] ; then
		echo "Restoring boot sector backup"
		if [[ "1" != "$DEBUG" ]] ; then
			dcfldd if=$SOURCE_DIR/boot-sector.img of=$DISK_DEVICE  || exit 44
		fi
		
		echo "Restoring partition table backup"
		if [[ "1" != "$DEBUG" ]] ; then
			sfdisk $DISK_DEVICE < $SOURCE_DIR/partition-table.sfdisk || exit 45
		fi
		
		echo "Updating from partition information"
		if [[ "1" != "$DEBUG" ]] ; then
			hdparm -z $DISK_DEVICE
		fi
	fi
	
	ls $SOURCE_DIR/partition*.tar$COMPRESSION_SUFFIX > /dev/null 2>&1
	if [[ "$?" != "0" ]] ; then
		echo "There are no files with the suffix .tar$COMPRESSION_SUFFIX in directory $SOURCE_DIR"
		echo "Please check you are using the right compression algorithm (options could be -j -J or -z)"
		exit 6
	fi
	
	
	#For each partition, we restore the backup
	for partition_file in $(ls $SOURCE_DIR/partition*.tar$COMPRESSION_SUFFIX) ; do
		partition_dev=$(basename $partition_file | sed "s+partition-+$DISK_DEVICE+" |sed "s/\.tar$COMPRESSION_SUFFIX//")
		backup_basename=$(basename $partition_file | sed "s/\.tar$COMPRESSION_SUFFIX//")
		backup_filename="$backup_basename.tar$COMPRESSION_SUFFIX "
		partition_type_file=$SOURCE_DIR/$backup_basename.type
		
		echo "--- Preparing backup restoration of $partition_file into $partition_dev ---"
		
		if [[ "$KEEP_PARTITIONS" != "1" ]] ; then
			#First get the partition type
			partition_type=$(cat $partition_type_file)
			prepare_partition "$partition_type" $partition_dev $backup_basename
		fi
		
		restore_partition $partition_file $_REAL_TARGET_DEVICE $backup_basename
	done
}

restore_disk(){
	if [[ -d "$SOURCE_DIR" ]] ; then
		echo "Restoring disk image from $SOURCE_DIR to $DISK_DEVICE"
		echo ""
		echo "You will lose any data stored on $DISK_DEVICE"
		echo "Please make sure you have a backup"
		echo ""
		echo "ARE YOU SURE you want to erase any data on $DISK_DEVICE ? (yes/NO)"
		read answer
		while [[ "$answer" =~ [yY]$ ]] ; do
			echo "Please type yes or YES"
			read answer
		done
		
		if [[ "$answer" =~ yes ]] ; then
			do_restore_disk
		fi
	else
		echo "ERROR: $SOURCE_DIR does not exist"
		exit 4
	fi
}

check_override_file(){
	local target_path=$1
	answer=""
	if [[ -f $target_path ]] ; then
		echo "Are you sure you want to destroy all data in $target_path (all previously saved data will be losts !!!)"
		echo "(yes/No)"
		read answer
		while [[ "$answer" =~ [yY]$ ]] ; do
			echo "Please type yes or YES"
			read answer
		done
		
		if [[ "$answer" =~ yes ]] ; then
			return 0
		else
			return 1
		fi
	fi
}

backup_mounted_volume(){
	backup_basename=$1
	partition_dev=$2
	my_part_type=$3
	
	backup_filename="$backup_basename.tar$COMPRESSION_SUFFIX "
	echo "--- Creating backup for $partition_dev into $backup_filename ---"
	mount_dir_name="diskimg-$backup_basename"
	mountpoint_dir="/mnt/$mount_dir_name"
	test "1" = "$VERBOSE" && echo "Mounting $partition_dev to $mountpoint_dir"
	mkdir -p $mountpoint_dir
	mount $partition_dev $mountpoint_dir
	test "1" = "$VERBOSE" && echo "Backing up from $mountpoint_dir to $TARGET_DIR/$backup_filename"
	
	if [[ "$SKIP_TAR" != 1 ]] ; then
		if check_override_file $TARGET_DIR/$backup_filename ; then
			if [[ "1" != "$DEBUG" ]] ; then
				echo "tar --numeric-owner $COMPRESSION_ALGORITHM $VERBOSE_OPTION $EXCLUDE_STRING -pcf $TARGET_DIR/$backup_filename -C $mountpoint_dir ."
				tar --numeric-owner $COMPRESSION_ALGORITHM $VERBOSE_OPTION $EXCLUDE_STRING -pcf $TARGET_DIR/$backup_filename -C $mountpoint_dir .
			fi
		fi
	fi
	
	umount $partition_dev
	rmdir $mountpoint_dir
}

ask_password(){
	local partition_dev=$1
	local need_confirmation=$2
	local second_try="FAKE"
	while [[ "$second_try" != "$_luks_password" ]] ; do
		echo "Enter encrypted partition passphrase for $partition_dev:"
		read -s _luks_password
		
		if [[ "$need_confirmation" = "confirm" ]]; then
			echo "Please confirm passphrase:"
			read -s second_try
			
			if [[ "$second_try" != "$_luks_password" ]] ; then
				echo "ERROR: passphrases do not match ! Please try again"
			fi
		else
			$second_try="$_luks_password"
		fi
	done
	_luks_password=$second_try
}

open_luks_container(){
	local target_device=$1
	local backup_basename=$2
	if [ "$DEBUG" != "1" ] && [ "$FAKE_DEVICE" != "1" ] ; then
		cryptsetup luksOpen $target_device $backup_basename<<EOF
$_luks_password
EOF
	fi
	_UNCRYPTED_DEVICE=/dev/mapper/$backup_basename
}

ask_and_open_encrypted_part(){
	local partition_dev=$1
	local backup_basename=$2
	#Are we already opened ? in this case, 
	#we should see "crypt" string in the output of lsblk -l -n $partition_dev
	lsblk -l -n $partition_dev | grep crypt > /dev/null 2>&1
	rv=$?
	if [ "$rv" != "0"  ] ; then
		ask_password $partition_dev
		open_luks_container $partition_dev $backup_basename
	fi
}

ensure_crypt_part_is_opened(){
	partition_dev=$1
	backup_basename=$2
	ask_and_open_encrypted_part $partition_dev $backup_basename
	#Now we should get the uncrypted volume name (which is not necessarily the 
	#same as $backup_basename)
	current_vol_name=$(lsblk -l -n $partition_dev |grep crypt | awk '{print $1}')
	
	#Now get the partition type and returns it
	my_part_type=`lsblk -l -n -o NAME,FSTYPE $partition_dev | grep $backup_basename | awk '{print $2}'`
	_UNCRYPTED_TYPE=${my_part_type}
	_UNCRYPTED_VOL=/dev/mapper/$current_vol_name
}

backup_disk(){
	if [[ ! -d "$TARGET_DIR" ]] ; then
		echo "$TARGET_DIR does not exist"
		echo "Do you want to create it ? (y/N)"
		read answer
		if [[ "$answer" =~ [yY] ]] ; then
			mkdir -p $TARGET_DIR
		else
			echo "Aborting because target dir does not exist"
			exit 126
		fi
	fi
	
	if [[ -d "$TARGET_DIR" ]] ; then
		echo "Creating disk image from $DISK_DEVICE to $TARGET_DIR..."
		
		echo "Creating boot sector backup"
		test "1" != "$DEBUG" && echo dcfldd if=$DISK_DEVICE bs=1M count=1  of=$TARGET_DIR/boot-sector.img
		test "1" != "$DEBUG" && dcfldd if=$DISK_DEVICE bs=1M count=1  of=$TARGET_DIR/boot-sector.img
		
		echo "Creating partition table backup"
		test "1" != "$DEBUG" && echo sfdisk -d $DISK_DEVICE > $TARGET_DIR/partition-table.sfdisk
		test "1" != "$DEBUG" && sfdisk -d $DISK_DEVICE > $TARGET_DIR/partition-table.sfdisk
		
		create_exclude_string
		
		#For each partition, we make a backup
		for partition_dev in $(sfdisk -d $DISK_DEVICE | grep -F 'start=' | awk '{print $1}') ; do
			backup_basename=$(echo $partition_dev | sed "s+$DISK_DEVICE+partition-+")
			
			#Device to mount will be empty unless necessary (ie: if a crypt volume is to be backup-ed)
			device_to_mount=""
			#Partition type
			echo "Managing partition type"
			part_type=`lsblk -lno NAME,FSTYPE $partition_dev | grep $(basename $partition_dev) | awk '{print $2}'`
			case $part_type in
			crypto_LUKS)
				#We extra steps here:
				#First decipher the luks partition:
				ensure_crypt_part_is_opened $partition_dev $backup_basename
				#The inner partition and device may have changed
				device_to_mount=${_UNCRYPTED_VOL}
				backup_mounted_volume $backup_basename $device_to_mount ${_UNCRYPTED_TYPE}
				#Concatenates the part types
				part_type="$part_type ${_UNCRYPTED_TYPE}"
				;;
			*)
				backup_mounted_volume $backup_basename $partition_dev $part_type
				;;
			esac
			part_type_filename="$backup_basename.type"
			echo $part_type > $TARGET_DIR/$part_type_filename
		done
	fi
}

if [ -z "$TARGET_DIR" ] && [ -z "$SOURCE_DIR" ]
then
	echo "You must indicate an operation with -c or -r"
	exit 2
fi

if [ -z "$DISK_DEVICE" ]
then
	echo "You must indicate a disk device"
	exit 3
fi

if [[ $EUID -ne 0 ]]; then
   echo "This script cannot work if not super-user. Please run as root" 1>&2
   exit 1
fi


if [[ "$TARGET_DIR" != "" ]] ; then
	backup_disk
	exit
fi

if [[ "$SOURCE_DIR" != "" ]] ; then
	restore_disk
	exit
fi