#!/bin/bash

# Init
COL_HIGHLIGHT="\033[0;33m"
COL_RESET="\033[0m"
AUTOMATIC=0
INDIR="<none>"
OUTFILE="<none>"

# Check if every tool is available
command -v partclone.restore >/dev/null || { echo -e "\n*** Error: partclone.restore not available\n"; exit 23; }

# Options parsing
while getopts ":i:ao:" optname; do
    case "$optname" in
      "i")
        INDIR=$OPTARG
        ;;
      "o")
        OUTFILE=$OPTARG
        ;;
      "a")
        AUTOMATIC=1
        ;;
      "?")
        echo -e "\n*** Error: unknown option -$OPTARG\n"
	exit 1
        ;;
      ":")
        echo -e "\n*** Error: missing argument for option -$OPTARG\n"
	exit 2
        ;;
      *)
        echo -e "\n*** Error: cannot process options\n"
	exit 99
        ;;
    esac
done

if [[ $INDIR == "<none>" ]]; then
        echo -e "\n*** Error: no -i option specified\n"
	exit 3
fi

if [[ $OUTFILE == "<none>" ]]; then
        echo -e "\n*** Error: no -o option specified\n"
	exit 4
fi

if [[ ! -d $INDIR ]]; then
        echo -e "\n*** Error: directory $INDIR doesn't exist, specify a different path\n"
	exit 5
fi

if [[ ! -b $OUTFILE ]]; then
        echo -e "\n*** Error: $OUTFILE is not a valid block device\n"
	exit 6
fi

if [[ ! -f $INDIR/mbr.img ]]; then
	echo -e "\n*** Error: cannot find MBR image\n"
	exit 7
fi

if   [[ -f $INDIR/boot.img ]]; then
	INEXT="img"
	COMPRESSION=0
	COMPR_DESCR="uncompressed"
elif [[ -f $INDIR/boot.img.bz2 ]]; then
	INEXT="img.bz2"
	COMPRESSION=1
	COMPR_DESCR="compressed"
else
	echo -e "\n*** Error: cannot find boot image\n"
	exit 8
fi

if [[ ! -f $INDIR/system.$INEXT ]]; then
	echo -e "\n*** Error: cannot find system image\n"
	exit 9
fi

echo -ne "\n${COL_HIGHLIGHT}Ready to perform restore of RPi sd card $OUTFILE from $COMPR_DESCR image in ${INDIR}${COL_RESET}\n"

if [[ $AUTOMATIC == 0 ]]; then
	echo
	read -p "Do you want to continue? (y/n) "
	if [[ ! $REPLY =~ [yY] ]]; then
		echo -e "\n*** Aborting\n"
		exit 98
	fi
fi

echo -e "\n${COL_HIGHLIGHT}Unmounting $OUTFILE...${COL_RESET}\n"

umount ${OUTFILE}?

echo -e "\n${COL_HIGHLIGHT}Restoring MBR...${COL_RESET}\n"

dd bs=512 count=1 of=$OUTFILE if=$INDIR/mbr.img
sync
partprobe $OUTFILE

echo -e "\n${COL_HIGHLIGHT}Restoring boot partition...${COL_RESET}\n"

if [[ $COMPRESSION == 1 ]]; then
	bunzip2 -c $INDIR/boot.$INEXT | partclone.restore --source - --output ${OUTFILE}1
else
	partclone.restore --output ${OUTFILE}1 --source $INDIR/boot.$INEXT
fi

echo -e "\n${COL_HIGHLIGHT}Restoring system partition...${COL_RESET}\n"

if [[ $COMPRESSION == 1 ]]; then
	bunzip2 -c $INDIR/system.$INEXT | partclone.restore --source - --output ${OUTFILE}2
else
	partclone.restore --output ${OUTFILE}2 --source $INDIR/system.$INEXT
fi

echo -e "\n${COL_HIGHLIGHT}Setting up swap partition...${COL_RESET}\n"

mkswap ${OUTFILE}3
sync

echo -e "\n${COL_HIGHLIGHT}Everything is Ok.${COL_RESET}\n"
