#!/bin/bash

# Informational output w/ happy colors
debug () {
    [[ $DEBUG ]] && echo -e '\033[1;33mDEBUG ::\033[1;m ' $*
}

info () {
    echo -e '\033[1;34m::\033[1;m ' $*
}

die () {
    echo -e '\033[1;31mFATAL ::\033[1;m ' $*
    exit 1
}

# Source default config, and bail on bad syntax
CONFIG=/etc/squashfu.conf
source $CONFIG
[[ $? -gt 0 ]] && die "Syntax error in config file."

am_i_root () {
    [[ $UID -eq 0 ]] || die "Must be root!"
}

create_directory_structure () {
    cd "${BKUP_ROOT}"
    mkdir -p {ro,rw,bins/{1..7}}
}

create_new_seed () {
    # Create a new squashfs based on the contents of the union
    # It's okay if we make an empty squash, we'll tell the user
    # about it at the end and suggest --resquash
    debug "Making new squash seed $(basename $SEED)"
    [[ $1 == "replace" ]] && {
        mksquashfs "${BKUP_ROOT}/rw" "$SEED.replace" -b 65536;
    } || {
        mksquashfs "${BKUP_ROOT}/rw" "$SEED" -b 65536;
    }
}

# The meat and potatoes of this bad agent.
do_backup () {
    [[ -w "${BKUP_ROOT}" ]] ||
        die "Backup root is not accessible. Please check your setting in /etc/squashfu"

    # Sanitize our workspace.
    unmount_all

    # Ensure that we have directories to mount and write to
    create_directory_structure

    # We need a seed to mount, even if it's going to be empty
    [[ -f "$SEED" ]] || {
        debug "No seed found -- creating a new one...";
        create_new_seed;
    }

    # Mount the necessary bins all as ro
    mount_union $(( $(date +%u) + $MODIFIER ))

    # Do the backup! Do doo doo doooo
    run_rsync

    # Force flush of pseudo links on branches to ensure clean dismount
    # Our umount.aufs helper can do this for us, but it doesn't handle
    # rsync hangups very well.
    sleep 5
    auplink "${BKUP_ROOT}/rw" flush

    # Is there reason to resquash?
    [[ $(date +%u) -eq $RESQUASH_DAY || $RESQUASH_AFTER ]] && {
        create_new_seed replace
        [[ $KEEP_LAST_WEEK ]] && {
            save_old_tree;
        } || {
            rm "$SEED";
            rm -r "${BKUP_ROOT}/bins"
        }

        mv "$SEED.replace" "$SEED"
        create_directory_structure
    }

    [[ $REPORT ]] && print_usage full

    unmount_all
}

mount_seed () {
    debug "Mounting seed"
    # Mount the squashed seed, failing if we can't
    mount -o loop,ro "${SEED}" "${BKUP_ROOT}/ro" || {
        die FATAL: Error mounting $SEED;
    }
}

mount_union () {
    # Account for second arg from rollback function to mount at a different mountpoint
    [[ -n "$2" ]] && {
        MOUNT_POINT="$(readlink -f $2)";
        echo "$MOUNT_POINT" >> /tmp/squashfu.custom
    } || {
        MOUNT_POINT="${BKUP_ROOT}/rw";
    }

    # build mount string
    # mount the first bin as rw, the rest as ro, and the seed as ro
    branches="br=${BKUP_ROOT}/bins/$1=rw:"
    for i in $(seq $(($1 - 1)) -1 1); do
        branches="${branches}${BKUP_ROOT}/bins/${i}=ro:"
    done
    branches="${branches}${BKUP_ROOT}/ro=ro"

    # build and execute mount command
    debug "mount -t aufs none "${MOUNT_POINT}" -o udba=reval,$branches"
    mount -t aufs none "${MOUNT_POINT}" -o udba=reval,$branches
}

print_usage () {
    # Do another sanity check -- check sizes of bins versus squash.
    total_bin_size=$(du -s ${BKUP_ROOT}/bins 2>/dev/null | awk '{print $1}')
    sfs_size=$(stat -c %s $SEED 2>/dev/null)

    # Basic query
    info "Overall Usage:"
    printf "%25s %10.2f MiB\n" "Seed size:" "$(echo "scale=2;$sfs_size / 1024 / 1024" | bc)"
    printf "%25s %10.2f MiB\n" "Total incremental size:" "$(echo "scale=2;$total_bin_size / 1024 / 1024" | bc)"

    # Full query
    [[ $1 == "full" ]] && {
        echo
        info "Detailed Usage By Bin:"
        for bin in {1..7}; do
            bin_size=$(du -s "${BKUP_ROOT}/bins/$bin" 2>/dev/null | awk '{print $1}');
            printf "%25s %10.2f MiB\n" "Bin $bin" $(echo "scale=2;$bin_size / 1024 / 1024" | bc);
        done;
    }

    # Compare seed and bin size. Suggest resquash if needed.
    [[ $total_bin_size -gt $sfs_size ]] && {
        info "Your incrementals are larger than your seed! You might consider resquashing your backup with $0 --resquash";
    }
}

run_rsync() {
    # Gather includes and excludes from config file
    # No error checking here -- user better not have
    # effed up the config
    INCLUDES=($(sed -n '/^<<INCLUDES$/,/^INCLUDES$/p' $CONFIG | grep -vE "^<*INCLUDES$"))
    EXCLUDES=($(sed -n '/^<<EXCLUDES$/,/^EXCLUDES$/p' $CONFIG | \
                grep -vE "^<*EXCLUDES$" | \
                sed -n 's/\(.*\)/--exclude "\1"/p'))

    # rsync source to $BKUP_ROOT/rw
    debug "Rsync executing with:"
    debug "   Options: ${RSYNC_OPTS[@]}"
    debug "   Includes: ${INCLUDES[@]}"
    debug "   Excludes: ${EXCLUDES[@]}"
    rsync ${RSYNC_OPTS[@]} ${INCLUDES[@]} ${EXCLUDES[@]} ${BKUP_ROOT}/rw

}

save_old_tree () {
    # create new directory, and then set aside old seed -- yes, you can do this while its mounted
    mkdir "${BKUP_ROOT}/last-week"
    cd "$BKUP_ROOT" && mv {$SEED,bins/} "${BKUP_ROOT}/last-week"
}

unmount_all () {
    am_i_root

    #Union must be unmounted first, or bad things happen
    unmount_custom
    unmount_union
    #unmount_seed
}

unmount_custom () {
    # Unmount any custom mounts from rollback operations
    debug "Checking for and unmounting user-specified mount points"
    [[ -f /tmp/squashfu.custom ]] && {
        while read mount; do
            umount $mount
        done < /tmp/squashfu.custom;
        rm /tmp/squashfu.custom;
    }
}

unmount_seed () {
    # Account for possibility of multiple mounts
    debug "Checking for and unmounting seed"
    while [[ $(mountpoint "${BKUP_ROOT}/ro" | grep "is a mount") ]]; do
        umount "${SEED}"
    done
}

unmount_union () {
    # Account for possibility of multiple mounts
    debug "Checking for and unmounting union"
    while [[ $(mountpoint "${BKUP_ROOT}/rw" | grep "is a mount") ]]; do
        debug "Found mount, attempting to unmount..."
        umount "${BKUP_ROOT}/rw"
    done
}

usage () {
    info "SquashFu: a backup solution hewn out of boredom"
    cat <<HELP

USAGE
    squashfu <operation> [options]

OPERATIONS
  -B
    Runs a regular backup, using the config file at /etc/squashfu, unless
    otherwise specified with the -c option.

  -R <day> [mount]
    Rollback to the day described by the following argument. 'day' needs to be
    a day of the week between the last resquash and the current day. An
    alternate mount point can be specified for the resulting rolled back union.
    If unspecified, the union will be mounted at $BKUP_ROOT/rw.

  -U
    Displays the size of the seed, the incrementals, and the actual backup. If
    you provide no additional options, a basic report will be given. Specifying
    "full" will give more detail about individual bins.

  -Z
    Use this option to restore sanity, AKA check for the seed or union being
    mounted and unmount them. Squashfu will ensure that everything is unmounted
    before a backup, but you might use this after restoring a file after a
    rollback.

OPTIONS
  -c, --config
    Specify an alternate location to a config file that will override defaults
    provided in /etc/squashfu.conf.

  --report
    Provide a usage report after the backup has completed.

  --resquash
    To be used with the -B operation. This forces a new squashed seed to be 
    created when the backup finishes. The original seed and incrementals are
    moved to the directory specified in the config file. Ensure that you have
    sufficient space for this operation, or bad things will happen. You can
    use the -Q operation to estimate how much free space you will need.

  --resquash-discard-old
    Similar to --resquash except the old seed and incrementals are discarded
    after the new seed is created.

HELP
    exit 0
}

#####################
#     Dispatch      #
#####################
[[ $# -eq 0 ]] && usage

dispatch_backup () {
    am_i_root
    while [[ $# -gt 0 ]]; do
        case $1 in
            "-c"|"--config")                                                     ;;
            "--report") REPORT=true                                              ;;
            "--resquash") RESQUASH_AFTER=true;KEEP_LAST_WEEK=true                ;;
            "--resquash-discard-old") RESQUASH_AFTER=true;KEEP_LAST_WEEK=false   ;;
            "--modifier") shift;MODIFIER=$1                                      ;;
            *) die "Invalid backup option $1"; usage                             ;;
        esac
        shift
    done

    do_backup

    exit 0
}

dispatch_report () {
    #am_i_root

    # Check for valid arg
    case $1 in
        "")     print_usage ;;
        "full") print_usage full ;;
        *)      die "Invalid report option"; usage ;;
    esac;

    exit 0
}

dispatch_rollback () {
    am_i_root
    # Check arguments conform
    [[ $# -eq 0 || $# -gt 2 ]] && {
        die "Invalid rollback option";
        usage;
    }

    # Ensure first rollback arg is a valid day of the week
    debug "Attempting to rollback to $1"
    date --date=$1 >/dev/null 2>&1
    [[ $? -gt 0 ]] && {
        die "Invalid day of week '$1'";
        usage;
    }

    # TODO: Make sure that user isn't trying to mount in the future?
    # This may not be an issue since the bins are cleaned after each
    # resquash.

    # Ensure second arg (if supplied) is a valid directory
    [[ -n "$2" && ! -d "$2" ]] && {
        die "'$2' is not a valid mount point"; 
        usage;
    }

    # Don't mount union multiple times
    unmount_custom
    unmount_union

    mount_seed

    # Convert day to numerical day of week and mount
    mount_union $(date --date=$1 +%u) $2

    exit 0
}

# Determine operation and send to appropriate dispatcher
while [[ $# -gt 0 ]]; do
    case $1 in
        "-B")  shift; dispatch_backup $*             ;;
        "-U")  shift; dispatch_report $1             ;;
        "-R")  shift; dispatch_rollback $*           ;;
        "-Z")  unmount_all                           ;;
        *)     usage                                 ;;
    esac
    shift
done