#!/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 () {
    ################################
    # Sanity checks
    ################################
    #  - Are we root?
    #  - is our BKUP_ROOT valid? (FAIL)
    #  - Check for pre-existing mounts just in case (and unmount them)
    #  - do we have a proper (expected) directory structure in place?
    #       (Use cd to BKUP_ROOT to avoid issues with brace expansion in a quoted path)

    [[ -w "${BKUP_ROOT}" ]] ||
        die "Backup root is not accessible. Please check your setting in /etc/squashfu"

    unmount_all

    create_directory_structure


    ################################
    # Prep work
    ################################
    #  - does seed exist? (if not, our backup is creating the seed)
    #  - Prepare union mount with proper bins
    [[ -f "$SEED" ]] || {
        debug "No seed found -- creating a new one...";
        create_new_seed;
    }

    mount_union_ro $(( $(date +%u) + $MODIFIER ))

    mount_union_branch_rw $(( $(date +%u) + $MODIFIER ))


    ################################
    # Call for backup!
    ################################
    run_rsync


    ################################
    # Cleanup 
    ################################
    #  - Is this resquash day? If so, we need a new squash
    #  - If new squash creation fails, we're in trouble. (by default, keep previous week)
    #  - Call unmount_all when we're done
    #  - Print a report if it were requested
    [[ $(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
    }

    unmount_all

    [[ $REPORT ]] && query_usage full
}

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

    # build branch string
    branches="br="
    for i in $(seq $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
}

mount_union_branch_rw () {
    debug "Remount branch $1 as read-write"
    mount -o remount,mod:bins/$1=rw "${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"
}

query_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)"

    if [[ $total_bin_size -gt $sfs_size ]]; then
        info "Your incrementals are larger than your seed! You might consider resquashing your backup with $0 --resquash"
    fi

    # 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;
    }
}

run_rsync() {
    # Gather includes and excludes from config file
    # No error checking here -- user better not have
    # effed up the config
    INCLUDES=($(grep ^#+ $CONFIG | cut -d+ -f2-))
    EXCLUDES=($(grep ^#- $CONFIG | cut -d- -f2-))

    # 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
}

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.

  -Q
    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.

  -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.

  -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.

  --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 option $1"; usage                                    ;;
        esac
        shift
    done

    do_backup
}

dispatch_query () {
    am_i_root
    # If no arg supplied, no need to parse
    [[ -z $1 ]] && query_usage

    # Check for valid arg
    case $1 in
        "full") query_usage full ;;
        *)      die "Invalid query option"; usage ;;
    esac
}

dispatch_rollback () {
    am_i_root
    # Check arguments conform
    [[ $# -eq 0 || $# -gt 2 ]] && {
        die "Invalid arguments to -R";
        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_ro $(date --date=$1 +%u) $2
}

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