#!/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:DhjJPr: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=""

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
        	;;
        -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

ESCAPED_COMPRESSION_SUFFIX="\$COMPRESSION_SUFFIX"

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


do_restore_disk(){
	echo "Restoring disk image from $SOURCE_DIR to $DISK_DEVICE"
	if [[ "$KEEP_PARTITIONS" != "1" ]] ; then
		echo "Restoring boot sector backup"
		test "1" != "$DEBUG" && dcfldd if=$SOURCE_DIR/boot-sector.img of=$DISK_DEVICE  || exit 44
		
		echo "Restoring partition table backup"
		test "1" != "$DEBUG" && sfdisk $DISK_DEVICE < $SOURCE_DIR/partition-table.sfdisk || exit 45
		
		echo "Updating from partition information"
		test "1" != "$DEBUG" && hdparm -z $DISK_DEVICE
	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$ESCAPED_COMPRESSION_SUFFIX//")
		backup_basename=$(basename $partition_file)
		backup_filename="$backup_basename.tar$COMPRESSION_SUFFIX "
		partition_type_file=$SOURCE_DIR/$(echo $backup_basename | sed "s+\.tar$ESCAPED_COMPRESSION_SUFFIX+.type+")
		
		echo ""
		echo "--> Restoring backup $partition_file into $partition_dev"
		
		if [[ "$KEEP_PARTITIONS" != "1" ]] ; then
			#First get the partition type
			partition_type=$(cat $partition_type_file)
			echo "Formating $partition_dev with type $partition_type"
			test "1" != "$DEBUG" && mkfs -t $partition_type $partition_dev
			test "$?" == 0 || exit $?
		fi
		
		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
		test "$?" == 0 || exit $?
		mount $partition_dev $mountpoint_dir
		test "$?" == 0 || exit $?
		test "1" = "$VERBOSE" && echo "Restoring from $partition_file to $mountpoint_dir"
		
		if [[ "$SKIP_TAR" != 1 ]] ; then
			test "1" = "$DEBUG" && echo tar --numeric-owner $COMPRESSION_ALGORITHM $VERBOSE_OPTION -pxf $partition_file -C $mountpoint_dir
			test "1" != "$DEBUG" && tar --numeric-owner $COMPRESSION_ALGORITHM $VERBOSE_OPTION -pxf $partition_file -C $mountpoint_dir
		fi
		
		test "$?" == 0 || exit $?
		test "1" = "$VERBOSE" && echo "Un-mounting $partition_dev"
		umount $partition_dev
		test "$?" == 0 || exit $?
		rmdir $mountpoint_dir
		test "$?" == 0 || exit $?
	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
}

backup_mounted_volume(){
	backup_basename=$1
	partition_dev=$2
	part_type=$3
	
	part_type_filename="$backup_basename.type"
	echo $part_type > $TARGET_DIR/$part_type_filename
	
	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
		test "1" != "$DEBUG" && echo "tar --numeric-owner $COMPRESSION_ALGORITHM $VERBOSE_OPTION $EXCLUDE_STRING -pcf $TARGET_DIR/$backup_filename -C $mountpoint_dir ."
		test "1" != "$DEBUG" && tar --numeric-owner $COMPRESSION_ALGORITHM $VERBOSE_OPTION $EXCLUDE_STRING -pcf $TARGET_DIR/$backup_filename -C $mountpoint_dir .
	fi
	
	umount $partition_dev
	rmdir $mountpoint_dir
}

ensure_crypt_part_is_opened(){
	partition_dev=$1
	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
		echo '### Please open the crypted volume by providing a passphrase to unlock it ###'
		cryptsetup luksOpen $partition_dev $backup_basename
	fi
	#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
	part_type=`lsblk -l -n -o NAME,FSTYPE $partition_dev | grep $backup_basename | awk '{print $2}'`
	_UNCRYPTED_TYPE=${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 "Saving partition type"
			part_type=`lsblk -lno NAME,FSTYPE $partition_dev | grep $(basename $partition_dev) | awk '{print $2}'`
			#part_type=`lsblk -lno NAME,FSTYPE $partition_dev | grep $DISK_DEVICE | print '{print $2}'`
			case $part_type in
			crypto_LUKS)
				#We extra steps here:
				#First decipher the luks partition:
				echo '###'$part_type'###'
				ensure_crypt_part_is_opened $partition_dev $backup_basename
				part_type=${_UNCRYPTED_TYPE}
				echo '***'$part_type'***'
				device_to_mount=${_UNCRYPTED_VOL}
				backup_mounted_volume $backup_basename $device_to_mount $part_type
				;;
			*)
				backup_mounted_volume $backup_basename $partition_dev $part_type
				;;
			esac
		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