wt-image

#!/bin/sh
# Copyright (C) 2013-2016 Francois Gouget
#
# Creates, checks, restores or reboots to a disk image.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
name0=`basename "$0"`

# Try to set a nicer output for the time command
time=""
time true >/dev/null 2>&1 && time="time"
TIME="Took %E (user %U system %S)"
export TIME


### Generic helper functions

error()
{
    echo "$name0:error: " "$@" >&2
}


opt_dry_run=""
dry_run()
{
    if [ -n "$opt_dry_run" ]
    then
        echo "$@" >&2
    else
        "$@"
    fi
}

img_wrapper()
{
    cmd="$1"
    filename="$2"
    case "$cmd" in
        check)
            if [ ! -f "$filename" -o ! -r "$filename" ]
            then
                error "'$filename' does not exist or is not readable"
                return 1
            fi
            ;;
        compress)
            dry_run $time cat >"$filename"
            ;;
        decompress)
            dry_run $time cat "$filename"
            ;;
    esac
}

gzip_wrapper()
{
    cmd="$1"
    filename="$2"
    _tool=`which pigz || which gzip`
    case "$cmd" in
        check)
            if [ ! -f "$filename" -o ! -r "$filename" ]
            then
                error "'$filename' does not exist or is not readable"
                return 1
            elif [ -z "$_tool" ]
            then
                error "could not find an xz tool for '$filename'"
                return 1
            fi
            ;;
        compress)
            dry_run $time "$_tool" -c >"$filename"
            ;;
        decompress)
            dry_run $time "$_tool" -c -d "$filename"
            ;;
    esac
}

bzip2_wrapper()
{
    cmd="$1"
    filename="$2"
    _tool=`which pbzip2 || which bzip2`
    case "$cmd" in
        check)
            if [ ! -f "$filename" -o ! -r "$filename" ]
            then
                error "'$filename' does not exist or is not readable"
                return 1
            elif [ -z "$_tool" ]
            then
                error "could not find an xz tool for '$filename'"
                return 1
            fi
            ;;
        compress)
            dry_run $time "$_tool" -c >"$filename"
            ;;
        decompress)
            dry_run $time "$_tool" -c -d "$filename"
            ;;
    esac
}

xz_wrapper()
{
    cmd="$1"
    filename="$2"
    _tool=`which pxz || which xz`
    case "$cmd" in
        check)
            if [ ! -f "$filename" -o ! -r "$filename" ]
            then
                error "'$filename' does not exist or is not readable"
                return 1
            elif [ -z "$_tool" ]
            then
                error "could not find a gzip tool for '$filename'"
                return 1
            fi
            ;;
        compress)
            dry_run $time "$_tool" -c >"$filename"
            ;;
        decompress)
            dry_run $time "$_tool" -c -d "$filename"
            ;;
    esac
}

borg_wrapper()
{
    cmd="$1"
    filename="$2"
    _tool=`which borg`
    BORG_UNKNOWN_UNENCRYPTED_REPO_ACCESS_IS_OK=yes
    export BORG_UNKNOWN_UNENCRYPTED_REPO_ACCESS_IS_OK
    case "$cmd" in
        check)
            if [ ! -d "$borg_repo" ]
            then
                error "'$borg_repo' is not a borg repository"
                return 1
            elif [ -z "$_tool" ]
            then
                error "could not find the borg backup binary"
                return 1
            elif "$_tool" list --list-format "{path}{NEWLINE}" "$borg_repo::$filename" 2>/dev/null | egrep "^(stdin|$filename\\.img)$" >/dev/null
            then
                return 0
            fi
            error "the '$borg_repo::$filename' borg archive does not exist or does not contain the expected file"
            return 1
            ;;
        compress)
            dry_run $time "$_tool" create -C zstd "$borg_repo::$filename" - "$wtimage_file"
            ;;
        decompress)
            borg_image=`borg list --list-format "{path}{NEWLINE}" "$borg_repo::$filename" 2>/dev/null | egrep "^(stdin|$filename\\.img)$"`
            dry_run $time "$_tool" extract --stdout "$borg_repo::$filename" "$borg_image"
            ;;
    esac
}


### Main

check_opt_val()
{
    option="$1"
    var="$2"
    argc="$3"

    if [ -n "$var" ]
    then
        error "$option can only be specified once"
        opt_usage=2 # but continue processing options
    fi
    if [ $argc -eq 0 ]
    then
        error "missing value for $option"
        opt_usage=2
        return 1
    fi
    return 0
}

opt_image=""
opt_backup=""
opt_description=""
opt_config=""
opt_action=""
opt_quiet=""
opt_usage=""
while [ $# -gt 0 ]
do
    arg="$1"
    shift
    case "$arg" in
    -n|--dry-run)
        opt_dry_run="1"
        ;;
   --check)
        if [ -n "$opt_action" ]
        then
            error "only one action can be specified"
            opt_usage=2
        else
            opt_action="check"
        fi
        ;;
   --grub)
        if [ -n "$opt_action" ]
        then
            error "only one action can be specified"
            opt_usage=2
        else
            opt_action="grub"
        fi
        ;;
   --restore)
        if [ -n "$opt_action" ]
        then
            error "only one action can be specified"
            opt_usage=2
        else
            opt_action="restore"
        fi
        ;;
   --backup)
        if [ -n "$opt_action" ]
        then
            error "only one action can be specified"
            opt_usage=2
        else
            opt_action="backup"
        fi
        if check_opt_val "$arg" "$opt_backup" $#
        then
            opt_backup="$1"
            shift
        fi
        ;;
   --description)
        if check_opt_val "$arg" "$opt_description" $#
        then
            opt_description="$1"
            shift
        fi
        ;;
   --config)
        if check_opt_val "$arg" "$opt_config" $#
        then
            opt_config="$1"
            shift
        fi
        ;;
    --quiet)
        opt_quiet="1"
        ;;
    --help|-h|-\?)
        opt_usage=0
        break
        ;;
    -*)
        error "unknown option '$arg'"
        opt_usage=2
        ;;
    *)
        if [ -z "$opt_image" ]
        then
            opt_image="$arg"
        else
            error "an image has already been specified"
            opt_usage=2
        fi
        ;;
    esac
done

if [ -z "$opt_usage" ]
then
    [ -n "$GZIP" ] || GZIP="-9"
    [ -n "$XZ_OPT" ] || XZ_OPT="-9"
    if [ -h "$opt_image" ]
    then
	opt_image=`readlink "$opt_image"`
    fi
    case "$opt_image" in
        *.img)
            wtimage_file=`echo "$opt_image" | sed -e 's/\.img$/.wtimage/'`
            compress="img_wrapper"
            ;;
        *.img.gz)
            wtimage_file=`echo "$opt_image" | sed -e 's/\.img\.gz$/.wtimage/'`
            compress="gzip_wrapper"
            ;;
        *.img.bz2)
            wtimage_file=`echo "$opt_image" | sed -e 's/\.img\.bz2$/.wtimage/'`
            compress="bzip2_wrapper"
            ;;
        *.img.xz)
            wtimage_file=`echo "$opt_image" | sed -e 's/\.img\.xz$/.wtimage/'`
            compress="xz_wrapper"
            ;;
        *.wtimage)
            wtimage_file="$opt_image"
            opt_image=`echo "$opt_image" | sed -e 's/\.wtimage$//'`
            if [ -f "$opt_image.img" ]
            then
                opt_image="$opt_image.img"
                compress="img_wrapper"
            elif [ -f "$opt_image.img.gz" ]
            then
                opt_image="$opt_image.img.gz"
                compress="gzip_wrapper"
            elif [ -f "$opt_image.img.bz2" ]
            then
                opt_image="$opt_image.img.bz2"
                compress="bzip2_wrapper"
            elif [ -f "$opt_image.img.xz" ]
            then
                opt_image="$opt_image.img.xz"
                compress="xz_wrapper"
            else
                hostname=`hostname | cut -d. -f1`
                borg_repo=`dirname "$opt_image"`
                if [ -d "$borg_repo/$name0.borg" ]
                then
                    borg_repo="$borg_repo/$name0.borg"
                    compress="borg_wrapper"
                    opt_image=`basename "$opt_image"`
                elif [ -d "$borg_repo/$hostname-$name0.borg" ]
                then
                    borg_repo="$borg_repo/$hostname-$name0.borg"
                    compress="borg_wrapper"
                    opt_image=`basename "$opt_image"`
                else
                    opt_image="$opt_image.img.xz"
                    compress="xz_wrapper"
                fi
            fi
            ;;
        *)
            if [ -n "$opt_image" ]
            then
                error "unknown image file type '$opt_image'"
                opt_usage=2
            fi
            ;;
    esac
    if [ -z "$opt_action" ]
    then
        error "you must specify the action to perform"
        opt_usage=2
    elif [ "$opt_action" = "grub" ]
    then
        # Nothing to do
        true
    elif [ "$opt_action" != "backup" ] && ! $compress check "$opt_image"
    then
        opt_usage=2
    elif [ -z "$opt_image" ]
    then
        error "you must specify the ntfs image to work on"
        opt_usage=2
    fi
    if [ -n "$opt_description$opt_config" -a "$opt_action" != "backup" ]
    then
        error "--description and --config can only be used with the --backup action"
        opt_usage=2
    fi
fi
if [ -n "$opt_usage" ]
then
    if [ "$opt_usage" != "0" ]
    then
        error "try '$name0 --help' for more information"
        exit $opt_usage
    fi
    echo "Usage: $name0 [--dry-run] [--check|--restore|--grub|--backup UUID]"
    echo "                [--description DESC] [--config CONFIG] IMAGE"
    echo

    echo "Creates, checks, restores or reboots to a disk image."
    echo

    echo "Where:"
    echo "  IMAGE         The disk image file or .wtimage file to work on. Disk images can"
    echo "                be compressed with gzip, bzip2, xz, or uncompressed."
    echo "                If a $name0.borg folder is in the same directory as the"
    echo "                .wtimage file, the disk image will be stored in that borg"
    echo "                repository."
    echo "  --dry-run     Show what would happen but do nothing."
    echo "  --check       Check if the partition still matches the disk image."
    echo "  --restore     Restores the disk image to the partition."
    echo "  --grub        Tells grub to reboot to the disk image's partition once."
    echo "  --backup UUID Creates a disk image of the specified partition. The UUID is"
    echo "                the partition's unique identifier."
    echo "  --description DESC A description to store in the .wtimage file for"
    echo "                documentation purposes when backing it up."
    echo "  --config CONFIG Extra configuration information to store in the .wtimage file"
    echo "                for documentation purposes when backing it up."
    echo "  --quiet       Silences some messages, particularly information about"
    echo "                the partition."
    echo "  --help, -h    Shows this help message."
    exit 0
fi

tmpdir="$TMPDIR"
[ -n "$tmpdir" ] || tmpdir="/tmp"
rcfile="$tmpdir/$name0.$$"


### Image file handling

read_wtimage_file()
{
    _ignore_missing=""
    if [ "$1" = "--ignore-missing" ]
    then
        _ignore_missing="1"
        shift
    fi
    _wtimage_file="$1"
    if [ -f "$_wtimage_file" ]
    then
        echo "$_wtimage_file" | grep / >/dev/null
        [ $? -eq 0 ] || _wtimage_file="./$_wtimage_file"
        . "$_wtimage_file"
        if [ -z "$img_tool" ]
        then
            error "this image was created using an unknown tool '$img_tool'"
            exit 1
        fi
        if [ -z "$img_uuid" ]
        then
            error "the wtimage file does not define the image uuid"
            exit 1
        fi
        if [ -z "$img_meta_sha1" ]
        then
            error "the wtimage file does not define the image's metadata sha1"
            exit 1
        fi
    elif [ -z "$_ignore_missing" ]
    then
        error "unable to find the wtimage file"
        exit 1
    fi
}

echo_wtimage()
{
    _uuid="$1"
    _sha1="$2"
    _description="$3"
    _config="$4"
    _wtimage=`perl -e "use Cwd; print Cwd::realpath('$0')"`
    echo "#!$_wtimage"
    echo "img_tool=ntfsclone"
    echo "img_uuid=$_uuid"
    echo "img_meta_sha1=$_sha1"
    echo "img_description=\"$_description\""
    echo "img_config=\"$_config\""
}

uuid_to_device()
{
    _uuid="$1"
    device="/dev/disk/by-uuid/$_uuid"
    if [ -h "$device" ]
    then
        device=`dirname "$device"`/`readlink "$device"`
    fi
    if [ ! -b "$device" ]
    then
        error "could not find a device matching '$_uuid'"
        exit 1
    fi
}

sha1_wrapper()
{
    if which sha1sum >/dev/null 2>&1
    then
        sha1sum "$@" | cut -d' ' -f1
    elif which sha1 >/dev/null 2>&1
    then
        sha1 "$@" | sed -e 's/^.* = //'
    fi
}

get_meta_sha1()
{
    device="$1"
    sha1=`(dry_run ntfsclone --metadata --save-image --output - "$device" 2>/dev/null; echo $? >"$rcfile") | sha1_wrapper`
    rc=`cat "$rcfile"`
    rm -f "$rcfile"
    [ $rc -ne 0 ] && return $rc
    echo "$sha1"
}

if [ "$opt_action" = "grub" ]
then
    [ -n "$img_label" ] || img_label=`sed -e "s/^menuentry ['\"]\(Windows[^'\"]*\)['\"].*\$/\\1/" -e t -e d /boot/grub/grub.cfg`
    if [ -z "$img_label" ]
    then
        error "unable to find the image's grub boot entry"
        exit 1
    fi
    grub-reboot "$img_label"

elif [ "$opt_action" = "backup" ]
then
    # Backup
    uuid_to_device "$opt_backup"
    read_wtimage_file --ignore-missing "$wtimage_file"
    if [ -n "$img_uuid" -a "$img_uuid" != "$opt_backup" ]
    then
        error "the '$opt_image' file is for partition '$image_meta_sha1', not '$opt_backup'"
        exit 1
    fi

    meta_sha1=`get_meta_sha1 "$device"`
    if [ $? -ne 0 ]
    then
        error "could not get the '$device' metadata checksum"
        exit 1
    fi
    if [ -z "$opt_dry_run" ]
    then
        echo_wtimage "$opt_backup" "$meta_sha1" "$opt_description" "$opt_config" >"$wtimage_file"
        chmod +x "$wtimage_file"
    else
        echo_wtimage "$opt_backup" "$meta_sha1" "$opt_description" "$opt_config"
    fi
    quiet=""
    [ -t 1 -o -n "$opt_quiet" ] || quiet="-q"
    ( dry_run ntfsclone $quiet --save-image "$device" --output -; echo $? >"$rcfile" ) |
        $compress compress "$opt_image"
    if [ $? -ne 0 ]
    then
        error "an error occurred while compressing to '$opt_image'"
        rm -f "$rcfile"
        exit 1
    fi
    rc=`cat "$rcfile"`
    rm -f "$rcfile"
    if [ $rc -ne 0 ]
    then
        error "could not save '$device'"
        exit 1
    fi

else
    # Check / Restore
    read_wtimage_file "$wtimage_file"
    device="/dev/disk/by-uuid/$img_uuid"
    if [ ! -b "$device" ]
    then
        error "could not find a device matching '$img_uuid'"
        exit 1
    fi
    meta_sha1=`get_meta_sha1 "$device"`
    if [ "$img_meta_sha1" != "$meta_sha1" ]
    then
        if [ "$opt_action" = "check" ]
        then
            if [ -z "$opt_quiet" ]
            then
                echo "Partition:   $device"
                echo "Description: $img_description"
                echo "Extra configuration: $img_config"
                echo "Part. sha1:  $meta_sha1"
                echo "Image sha1:  $img_meta_sha1"
                echo "The partition needs re-imaging"
            fi
            exit 3
        fi
        quiet=""
        [ -t 1 -o -n "$opt_quiet" ] || quiet="-q"
        ( $compress decompress "$opt_image"; echo $? >"$rcfile" ) |
            dry_run ntfsclone $quiet --restore-image --overwrite "$device" -
        if [ $? -ne 0 ]
        then
            error "could not restore '$device'"
            rm -f "$rcfile"
            exit 1
        fi
        rc=`cat "$rcfile"`
        rm -f "$rcfile"
        if [ $rc -ne 0 ]
        then
            error "could not decompress '$opt_image'"
            exit 1
        fi
    elif [ "$opt_action" = "check" ]
    then
        echo "The partition matches the image"
    fi
fi