#!/bin/bash

usage="Usage: $0 DEVICE DIR
Try \`$0 --help' for more information."

help="Usage: $0 DEVICE DIR

Restore the entire contents of disk drive DEVICE from the compressed
image files inside directory DIR. The disk is assumed to have 512-byte
sectors and DOS-style (MBR) partitioning, and partitions on DEVICE are
assumed to be accessible via their own device files named by appending
the partition number to DEVICE.

Example:
$0 /dev/sda /mnt0/sda-image/

Files in DIR must be named pN-START-SIZE-ID-TOOL-VERSION.gz for
partition images, or d-START-SIZE-ID-TOOL-VERSION.gz for images of
disk regions outside named partitions, where N is the partition
number, ID is the partition ID (type) in hex, START is the LBA within
DEVICE of the first sector copied to the image, SIZE is the number of
sectors the imaged region occupies, TOOL is the name of the imaging
tool used to create the file, and VERSION is that tool's version
number.  All files are compressed with gzip.

Author: Stephen Thomas <flabdablet@gmail.com> 15-Nov-2008
This is free software.  Do anything you like with it except
hold me accountable for any grief it causes you."

if [[ -z "$1" ]]
then
	echo "$usage"
	exit 2
fi

if [[ "$1" == "--help" || "$1" == "-h" ]]
then
	echo "$help"
	exit 2
fi

# Validate arguments

self="${0##*/}"
dev="$1"
dir="${2:-$(pwd)}"
dir="${dir%/}"

if [[ ! -e "$dir" ]]
then
	echo >&2 $self: $dir: not found
	exit 2
fi

if [[ ! -d "$dir" ]]
then
	echo >&2 $self: $dir: not a directory
	exit 2
fi

# Define the partition copy tools.  If multiple versions are known to be compatible,
# include all their version numbers in the same "compatible" string, separated and
# surrounded by spaces e.g. compatible[t]=' 1.13.1 1.13.2 1.13.3 '

t=0

tool[t]='dd'
compatible[t]=''
getversion[t]='dd --version | sed -n '\''1s/[^0-9.]*//p'\'
invoke[t++]='zcat "$source" | dd of="$dest" bs=$bs $seek conv=notrunc'

tool[t]='ntfsclone'
compatible[t]=' 1.12.1 '
getversion[t]='"$cmd" 2>&1 | sed -n '\''1s/[^0-9]*\([^ ]*\).*/\1/p'\'
invoke[t++]='zcat "$source" | "$cmd" --restore-image --overwrite "$dest" -'

tool[t]='ntfsclone'
compatible[t]=' 1.13.1 '
getversion[t]='"$cmd" 2>&1 | sed -n '\''1s/[^0-9]*\([^ ]*\).*/\1/p'\'
invoke[t++]='zcat "$source" | "$cmd" --restore-image --overwrite "$dest" -'

tool[t]='ntfsclone'
compatible[t]=' 2.0.0 '
getversion[t]='"$cmd" 2>&1 | sed -n '\''1s/[^0-9]*\([^ ]*\).*/\1/p'\'
invoke[t++]='zcat "$source" | "$cmd" --restore-image --overwrite "$dest" -'

tool[t]='partimage'
compatible[t]=' 0.6.4 '
getversion[t]='"$cmd" -v | sed -n '\''1s/[^0-9]*\([^ ]*\).*/\1/p'\'
invoke[t++]='"$cmd" -z1 -c -o -d -V -e -M -f3 -b -B "*=Continue;" restore "$dest" "$source"'

# Define a few utility functions

# Remember arguments as a deferred command line

function defer
{
	deferred_commands="$deferred_commands$*"$'\n'
}

# Quote argument list: allows items that may contain quotes
# to be passed safely to defer()

function quote
{
	printf '%q ' "$@"
}

# Given a filename, generate a deferred image-load script

function deferred_load
{
	# Validate image filename

	match='/([dp])([0-9]*)-([0-9]+)-([0-9]+)-([^-]+)-([^-]+)-([^-]+)\.gz$'
	if [[ "$1" =~ $match ]]
	then
		typeflag=${BASH_REMATCH[1]}
		partition=${BASH_REMATCH[2]}
		start=${BASH_REMATCH[3]}
		size=${BASH_REMATCH[4]}
		id=${BASH_REMATCH[5]}
		creator=${BASH_REMATCH[6]}
		version=${BASH_REMATCH[7]}
	else
		echo >&2 $self: $1: invalid image filename
		exit 2
	fi

	# Generate destination device name

	if [[ "$typeflag" == "d" ]]
	then
		partition=''
	fi
	dest="$dev$partition"

	# Choose a compatible image restoration tool
	# (same name, compatible version)

	tool_unknown=1
	for t in ${!tool[*]}
	do
		if [[ "${tool[t]}" == "$creator" &&
			-z "${compatible[t]}" ||
			"${compatible[t]}" == *" $version "* ]]
		then
			tool_unknown=0
			break
		fi
	done

	if ((tool_unknown))
	then
		echo >&2 "$self: Partition $partition needs $creator $version"
		echo >&2 "but I don't know how to use it.  Please edit me and"
		echo >&2 "add it to my list of partition copy tools."
		exit 2
	fi

	# Make sure a compatible version of the tool is available

	cmd_unavailable=1
	candidates="${tool[t]}-$version ${tool[t]}"
	alternatives=${compatible[t]/ $version / }
	for v in $alternatives
	do
		candidates="$candidates ${tool[t]}-$v"
	done
	for cmd in $candidates
	do
		if which $cmd >/dev/null 2>&1 &&
		[[ -z "${compatible[t]}" ||
		 "${compatible[t]}" == *" $(eval ${getversion[t]}) "* ]]
		then
			cmd_unavailable=0
			break
		fi
	done

	if ((cmd_unavailable))
	then
		echo >&2 "$self: Partition $partition needs $creator $version"
		echo >&2 "but I can't find it.  Please make it available as"
		echo >&2 "'$creator' or '$creator-$version'."
		if [[ "$alternatives" != ' ' ]]
		then
			echo >&2 "Compatible versions are OK:$alternatives"
		fi
		exit 2
	fi

	# Having found the right commands to use, generate deferred script

	defer cmd=$(quote "$cmd")
	defer source=$(quote "$1")
	defer dest=$(quote "$dest")
	if [[ "$typeflag" == "d" ]]
	then
		# When copying unpartitioned sectors to the
		# underlying device, do them one sector at
		# a time and put them at the correct offset

		defer bs=512
		defer seek=seek=$start
	else
		# When using dd to copy a partition, do the
		# whole thing in 1M chunks and don't force
		# an offset

		defer bs=1M
		defer seek=

		# If the target device is a real block-special disk device,
		# re-reading the partition table will make udev tear down and
		# rebuild all its subdevices, and might make hald try to mount
		# them. Make sure images get written to the proper targets, not
		# devices that will shortly disappear or (worse) accidentally
		# created regular files with the same names.

		if [[ -b "$dev" ]]
		then
			# If this is a device that existed before load-image
			# modified the partition table, test the inode-change
			# time as well as block-device existence, to avoid writing
			# to a subdevice that udev hasn't even had time to tear down

			ctime=unknown
			for ((i=0; i<devs; i++))
			do
				if [[ "$dest" == "${devname[i]}" ]]
				then
					ctime=${devtime[i]}
					break
				fi
			done

			# Generate script to wait for new device to exist

			defer "echo >&2 $self: waiting for new $dest"
			defer 'until'
			defer 'stat=($(stat -c'\''%F %Z %G'\'' $dest 2>/dev/null)) &&'
			defer '[[ "${stat[0]} ${stat[1]} ${stat[2]}" == "block special file" ]] &&'
			defer '[[ ${stat[3]} !=' $ctime ']]'
			defer 'do' 
			defer 'sleep 1'
			defer 'done'

			# Generate script to wait for device to mount and
			# then unmount it, if it belongs to plugdev group

			defer 'if [[ "${stat[4]}" == "plugdev" ]]'
			defer 'then'
			defer "echo >&2 $self: waiting for $dest to automount"
			defer 'tries=0'
			defer 'until mount | grep "^$dest " >/dev/null || ((++tries >= 15))'
			defer 'do'
			defer 'sleep 1'
			defer 'done'
			defer 'if mount | grep "^$dest " >/dev/null'
			defer 'then'
			defer "echo >&2 $self: unmounting $dest"
			defer 'tries=0'
			defer 'until umount "$dest" 2>/dev/null || ((++tries >= 15))'
			defer 'do'
			defer 'sleep 1'
			defer 'done'
			defer 'fi'
			defer 'fi'
		fi 
	fi
	defer "echo >&2 $self: writing partition ${partition:-"table at LBA $start"} to $dest using $cmd"
	defer "${invoke[t]}"
	defer '(($? && ++errors))'
}

# Initialize restoration script

deferred_commands=
defer 'errors=0'

# Generate unmount commands as needed for target device and all existing
# subdevices, and remember their names and inode-change timestamps

devs=0
for dest in "$dev"*
do
	if [[ -b "$dest" ]]
	then
		defer "if mount | grep $(quote "^$dest ") >/dev/null"
		defer 'then'
		defer "echo >&2 $self: unmounting $dest"
		defer 'tries=0'
		defer "until umount $(quote "$dest") 2>/dev/null || ((++tries >= 15))"
		defer 'do'
		defer 'sleep 1'
		defer 'done'
		defer 'fi'
		devname[devs]="$dest"
		devtime[devs++]=$(stat -c%Z "$dest")
	fi
done

# Generate restoration commands for all the disk-region images -
# these will rebuild the entire partition table, including those
# portions embedded in extended partitions

for image in "$dir"/d-*
do
	deferred_load "$image"
done

defer 'if ((!errors))'
defer 'then'

# If the target is a block-special device, generate script to
# re-read the partition table

if [[ -b "$dev" ]]
then
	defer "echo >&2 $self: re-reading partition table"
	defer 'sfdisk --re-read' $(quote "$dev")
fi

# Generate restoration commands for all the partition images

for image in "$dir"/p*
do
	deferred_load "$image"
done

defer 'fi'

# Final confirmation before work actually starts

echo >&2 $self: about to overwrite $dev with partition images from $dir.
read -p "Is this OK? (y/N) "
if [[ "${REPLY:0:1}" != y && "${REPLY:0:1}" != Y ]]
then
	exit 1
fi

# If we got this far, all sanity checks have passed and we have
# a complete image-loading script built up in $deferred_commands;
# run it

eval "$deferred_commands"

if ((errors))
then
	((errors != 1)) && plural=s
	echo >&2 "$self: $errors error$plural occurred while loading images."
	echo >&2 "$dev may be corrupt."
	exit 1
fi
echo >&2 "$self: successful completion."