#!/bin/bash

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

help="Usage: $0 DEVICE DIR

Make compressed images of all partitions and partition tables of the
disk drive DEVICE in files inside directory DIR.  The disk is assumed
to have 512-byte sectors and DOS-style (MBR) partitioning, and the
partitions are assumed to be accessible via their own device files
named by appending the partition number to DEVICE.  Partition imaging
tools are used to make images of recognized partition types, avoiding
access to unused disk blocks.

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

Files in DIR are 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 via gmail> 7-Oct-2008
This is free software.  Do anything you like with it except
blame me for its many shortcomings."

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

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

# Validate arguments

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

if [ -e "$dir" ]
then
	if [ -d "$dir" ]
	then
		echo >&2 $self: about to create partition images of $dev in $dir.
		read -p "Is this OK? (y/N) "
		if [ "${REPLY:0:1}" != y -a "${REPLY:0:1}" != Y ]
		then
			exit 1
		fi
	else
		echo >&2 $self: $dir: not a directory
		exit 1
	fi
else
	mkdir "$dir"
fi

# Generate a list of partitions to save, including extended partitions
# (plus a fake "extended partition" for the MBR+bootloader region)

{
	sfdisk -l -uS -x "$dev" || exit 1
} | {
	i=0
	partition[i]=
	start[i]=0
	size[i]=-1
	id[i]=5

	while read f1 f2 f3 f4 f5 f6 f7
	do
		if [[ "$f1" == "$dev"* ]]
		then
			if [ "$f2" == '*' ]
			then
				f2=$f3; f3=$f4; f4=$f5; f5=$f6
			fi
			if [ "$f5" != 0 ]
			then
				f1=${f1#-}
				partition[++i]=${f1#$dev}
				start[i]=$f2
				size[i]=$f4
				id[i]=$f5
			fi
		fi
	done

# Walk the list of partitions, trimming the sizes recorded for
# extended partitions so that they only cover blocks before the
# start of the next logical partition.  Also remove the partition
# number for extended partitions, so that the copy logic will
# directly access the underlying device instead of attempting to
# use the extended-partition container (the kernel has funny
# ideas about the sizes of devices mapped to extended partitions).

	extended=" 5 f 85 "
	for i in ${!partition[*]}
	do
		if [[ $extended =~ " ${id[i]} " ]]
		then
			for j in ${!partition[*]}
			do
				if [[ ! $extended =~ " ${id[j]} " ]]
				then
					((diff = start[j] - start[i]))
					if ((size[i] < 0 || diff >= 0 && diff < size[i]))
					then
						((size[i] = diff))
					fi
				fi
			done
			partition[i]=
		fi
	done

# Define the partition copy tools

	t=0

	tool[t]='dd'
	types[t]=''
	invoke[t]='dd if="$source" $bs $skip $count | gzip -9 >"$dest"'
	version[t++]="$(dd --version | sed -n '1s/[^0-9.]*//p')"

	tool[t]='ntfsclone'
	types[t]='7 17 27'
	invoke[t]='ntfsclone --save-image --output - "$source" | gzip -9 >"$dest"'
	version[t++]="$(ntfsclone 2>&1 | sed -n '1s/[^0-9]*\([^ ]*\).*/\1/p')"
	
	tool[t]='partimage'
	types[t]='4 6 b c e 12 14 16 1b 1c 1e 83 de'
	invoke[t]='partimage -z1 -c -o -d -V -M -f3 -b -B '\''*=Continue;'\'' save "$source" "$dest"'
	version[t++]="$(partimage -v | sed -n '1s/[^0-9]*\([^ ]*\).*/\1/p')"

	tool[t]='dd'
	types[t]='82'
	invoke[t]='dd if="$source" bs=4096 count=1 | gzip -9 >"$dest"'
	version[t++]="$(dd --version | sed -n '1s/[^0-9.]*//p')"

# Assign each partition a copy tool based on its partition ID

	for i in ${!partition[*]}
	do
		copywith[i]=0
		for t in ${!tool[*]}
		do
			if [[ " ${types[t]} " == *" ${id[i]} "* ]]
			then
				copywith[i]=$t
			fi
		done
	done

# Make the partition images

	for i in ${!partition[*]}
	do
		# Try with the selected tool, then
		# again with dd if a fancy tool fails

		t=${copywith[i]}
		until
			echo >&2
			echo >&2 $self: saving partition ${partition[i]:-"table at LBA ${start[i]}"} with ${tool[t]}:

			source=$dev${partition[i]}
			if [ -z ${partition[i]} ]
			then
				# When copying unpartitioned sectors from the
				# underlying device, get only those required

				bs="bs=512"
				skip="skip=${start[i]}"
				count="count=${size[i]}"
				dest="$dir/d-${start[i]}-${size[i]}-${id[i]}-${tool[t]}-${version[t]}.gz"
			else
				# When copying a partition, let dd do the
				# whole thing in 1M chunks

				bs="bs=1M"
				skip=
				count=
				dest="$dir/p${partition[i]}-${start[i]}-${size[i]}-${id[i]}-${tool[t]}-${version[t]}.gz"
			fi
			eval ${invoke[t]}
		do
			rm -f "$dest"
			if ((t))
			then			
				echo >&2 $self: copy failed - retrying with dd
				t=0
			else
				echo >&2 $self: copy failed - giving up
				exit 1
			fi
		done
	done
	echo >&2
	echo >&2 "$self: successful completion."
}