#!/bin/bash
VER="1.0"

CONFIG="/etc/squashfu.conf"
. "$CONFIG"

# Informational output w/ happy colors
debug () {
  [[ "$DEBUG" == "true" ]] && printf '\033[1;33mDEBUG ::\033[1;m %s\n' "$*"
}

warn () {
  printf '\033[1;33m::\033[1;m %s\n' "$*"
}

info () {
  printf '\033[1;34m::\033[1;m %s\n' "$*"
}

die () {
  printf '\033[1;31m::\033[1;m %s\n' "$*" >&2
  exit 1
}

create_new_squash () {
# Args: number of bins to be squashed, -1 on initial creation
# Returns: 0 on success, non-zero on failure

  # If making first seed, create it directly from source
  if [[ $1 -eq -1 ]]; then
    info "Creating seed from sources (this may take a while)"
    mksquashfs "${INCLUDES[@]}" "$SEED" -b 65536 -e "${EXCLUDES[@]}" >/dev/null
    if [[ $? -eq 0 ]]; then
      mount_squash
      info "Seed creation finished. It has been mounted at "$SQUASH_MOUNT" if you would like to make sure the proper files are included"
      return 0;
    else
      die "There was an error creating the initial squash."
    fi
  fi

  # Determine oldest $1 bins and mount them with the current squash
  local old_bins=($(sort -n -r -t: -k2 "$BINVENTORY" | tail -$1 | cut -d: -f1))

  mount_union_with_bins ${old_bins[@]}

  info "Merging old incrementals"
  # Create new squash with temp name, exiting on failure
  mksquashfs "$UNION_MOUNT" "$SEED.replace" -b 65536 >/dev/null || return 1

  unmount_all

  # Replace old squash
  mv "${SEED}.replace" "$SEED"

  info "Cleaning up inventory"
  # Delete old bins, and remove entry from binventory
  for bin in ${old_bins[@]}; do
    rm -rf "${BINS_DIR}/$bin"
    sed -i "/^$bin:/d" "$BINVENTORY"
  done
}

create_new_bin () {
# Arguments: 1, the number of the bin to create
# Returns: 0 on success, non-zero on error

  debug "Asked to create new bin: $1"
  # Create new directory, fail if it exists (something's wrong)
  mkdir "${BINS_DIR}/$1"
  [[ $? != 0 ]] && return $?

  # Update binventory with new bin name and timestamp
  echo "${1}:$(date +%s)" >> "$BINVENTORY"

  # If write to bin list fails, remove diretory and exit
  if [[ $? -ne 0 ]]; then
    rmdir "${BINS_DIR}/${1}"
    die "Error writing to '$BINVENTORY'"
  fi

  return
}

# Mounting functions
mount_squash () {
# Arguments: none
# Returns: return code of mount command
  mountpoint -q "$SQUASH_MOUNT" || {
    info "Mounting squash";
    mount -o loop,ro "$SEED" "$SQUASH_MOUNT";
  }
  return $?
}

mount_union_with_bins () {
# Arguments: numbers of bins to be mounted (variable number)
# Returns: 0 on successful mount, non-zero on failure
  info "Mounting union"
  debug "Requested to mount bins: $*"

  # Mount first as rw, shift, and mount the rest ro
  branches="br=${BINS_DIR}/$1=rw:"; shift
  if [[ -n $1 ]]; then
    for bin in $*; do
      branches="${branches}${BINS_DIR}/$bin=ro:"
    done
  fi
  branches="${branches}${SQUASH_MOUNT}=ro"

  mount -t aufs none "$UNION_MOUNT" -o udba=reval,$branches

  return $?
}

# Unmounting functions
unmount_union () {
# Args: none
# Returns: return code from umount
  info "Unmounting union"
 umount "$UNION_MOUNT" 2>/dev/null
 sleep .5
  return $?
}

unmount_squash () {
# Args: none
# Returns: return code from umount
  info "Unmounting squash"
  umount "$SQUASH_MOUNT" 2>/dev/null
  return $?
}

unmount_all () {
# Args: none
# Returns: none

  [[ $UID -ne 0 ]] && die "Must be root to unmount."

  # Union MUST be unmounted first
  mountpoint -q "$UNION_MOUNT" && unmount_union
  mountpoint -q "$SQUASH_MOUNT" && unmount_squash
}

check_for_resquash () {
# Args: none
# Returns: number of bins needing to be merged
  local number_of_bins=$(grep -v "^[[:space:]]*$" "$BINVENTORY" | wc -l)
  debug "Found $number_of_bins bins"

  if [[ $number_of_bins -gt $MAX_BINS ]]; then
    return $(( $number_of_bins - $MIN_BINS ))
  else
    return 0
  fi
}

get_next_available_bin () {
# Arguments: none
# Returns: Numeric value of the next unused bin
  for (( i=1; i <= MAX_BINS + 1; i++ )); do
    [[ -d "${BINS_DIR}/$i" || $(grep "^$i:" "$BINVENTORY") ]] && continue;
    debug "Next available bin = $i"
    return $i
  done
}

action_backup () {
# Args: options array squashfu was invoked with, shifted 1
# Returns: none

  [[ $UID -ne 0 ]] && die "Must be root to perform a backup"

  # Does the binventory exist? If not, prompt to make sure this is an initialization
  if [[ ! -f "$BINVENTORY" || ! -f "$SEED" ]]; then
    echo -ne "\033[1;33m::\033[0m"
    read -N1 -p "Looks like this is your first time running SquashFu. Is this correct? [y/N] " reply
    echo
    [[ ! "$reply" =~ [Yy] ]] && die "Your bin inventory and/or seed seem to be missing. Please fix this before continuing."

    # If we got here, the user answered yes, so initialize a new structure
    mkdir -p "$UNION_MOUNT" "$SQUASH_MOUNT" "${BINS_DIR}"
    touch "$BINVENTORY"
    create_new_squash -1
    exit 0
  fi

  info "Backup requested at $(date --rfc-3339=seconds)"

  # Cleanup union, in case user was doing a rollback and forgot to unmount (or error on last run)
  mountpoint -q "$UNION_MOUNT" && unmount_union

  info "Creating new bin"
  # Make a new bin for this incremenetal
  get_next_available_bin
  local new_bin=$?
  create_new_bin $new_bin

  # Determine the mount order via binventory
  local bin_order=($(sort -n -r -t: -k2 "$BINVENTORY" | cut -d: -f1))

  mountpoint -q "${SQUASH_MOUNT}" || mount_squash
  mount_union_with_bins ${bin_order[@]} || return 1

  # Includes are pulled in directly from config
  EXCLUDES=$(for excl in ${EXCLUDES[@]}; do echo --exclude $excl; done)

  debug "rsync ${RSYNC_OPTS[@]} ${INCLUDES[@]} ${EXCLUDES[@]} "$UNION_MOUNT""
  info "Creating new incremental"
  /usr/bin/rsync ${RSYNC_OPTS[@]} ${INCLUDES[@]} ${EXCLUDES[@]} "$UNION_MOUNT"
  rsync_ret=$?

  for error in ${DEL_BIN_ON_FAIL[@]}; do
    if [[ $rsync_ret == $error ]]; then
      warn "Unexpected hangup by rsync ($error). Deleting backup."
      action_remove_bin $new_bin override
      break
    fi
  done

  check_for_resquash
  [[ val=$? -gt 0 ]] && create_new_squash $val

  # TODO: Report if requested

  unmount_all

  info "Backup completed at $(date --rfc-3339=seconds)"
}

action_remove_bin () {
  # check if the bin exists both in the binventory AND in the bins directory
  [[ $UID != 0 ]] && die "Must be root to remove a backup"
  [[ ! -w "$BINVENTORY" ]] && die "Unable to write to ${BINVENTORY}"

  if grep "^$1:" ${BINVENTORY} && [[ -d "${BINS_DIR}/$1" ]]; then
    if [[ -z $2 ]]; then
      echo "Are you SURE you want to remove this bin?"
      local timestamp=$(sed -n "/^$1:/s/^[0-9]*:\([0-9]*\)/\1/p" "${BINVENTORY}")
      printf "\t%15s %s\n\t%15s %s\n\t%15s %s\n" \
          "Bin:" "$1" \
          "Date Created:" "$(date --rfc-3339=seconds --date="@$timestamp")" \
          "Size:" "$(du -sh "${BINS_DIR}/$1" 2>/dev/null | awk '{print $1}')"

      echo -ne "\033[1;33m::\033[0m"
      read -N1 -p "Confirm deletion (y/N) " reply
      echo
      [[ ! "$reply" =~ [Yy] ]] && die "Delete operation aborted"
    fi

    info "Deleting bin $1"
    sed -i "/^$1:[0-9]*/d" "${BINVENTORY}"
    rm -rf ${BINS_DIR}/$1
  else
   die "Bin $1 not found."
  fi
}

action_rollback () {
# Args: number of backups to roll back
# Returns: none

  [[ $UID -ne 0 ]] && die "Must be root to perform a rollback"
  [[ ! $1 =~ ^[0-9]*$ ]] && die "Invalid argument. Parameter must be a number."

  # Form a chronologically ordered list of bins, assuming the user didn't give bogus input
  local bin_list=($(sed -n 's/[\t\ ]*\([0-9]*\):.*/\1/p' "$BINVENTORY"))

  [[ $1 -gt ${#bin_list[@]} ]] && die "Cannot rollback more than ${#bin_list[@]} backups"

  local num_to_mount=$(( ${#bin_list[@]} - $1 ))

  mountpoint -q "$UNION_MOUNT" && unmount_union
  mountpoint -q "$SQUASH_MOUNT" || mount_squash

  if [[ ${#bin_list[@]} -eq $1 ]]; then
    local rb_timestamp=0 # XXX: What the hell is the timestamp of the seed... ?
  else
    mount_union_with_bins ${bin_list[@]:(-$num_to_mount)}
    local rb_timestamp=$(grep "^${bin_list[@]:(-$num_to_mount):1}:" "$BINVENTORY" | cut -d: -f2)
  fi

  info "You have rolled back to $(date --rfc-3339=seconds --date="@$rb_timestamp")"
  info "Your files can be found at '${UNION_MOUNT}'"
}

action_report () {
  info "SquashFu Usage Report"
  # Enumerate bins, sort date order, print human readable create date and size

  pushd "$BINS_DIR" &>/dev/null

  # Collect all data into an array to 'preload' it. Index 0 is the entire
  # folder. The following indicies correspond to the bin number of that index
  printf "\n%30s\r" ".: Loading :." >&2
  while read size bin; do
    case ${bin} in
      'total') total=$size; continue ;;
      '.') DATA[0]=$size ;;
      *) DATA[${bin}]=$size ;;
    esac
  done < <(du -csh . * 2>/dev/null | sort -n -k2)
  printf "%30s\r" "" >&2

  printf "%10s\t%25s\t%7s\n" "Bin ID" "Date Created" "Size"
  OIFS=$IFS;IFS=${IFS}$':'
  while read bin stamp; do
    printf "%10d\t%25s\t%7s\n" "$bin" "$(date --rfc-3339=seconds --date="@$stamp")" "${DATA[$bin]}"
  done < <(grep -v "^[[:space:]]*$" "$BINVENTORY" | sort -nr -t':' -k2)
  IFS=$OIFS
  printf "%10s\t%25s\t%7s\n" "" "Incremental Total" "${DATA[0]}"

  # Print totals
  printf "\n%10s\t%25s\t%7s\n" "" "$(basename $SEED)" "$(du -h "${SEED}" | awk '{print $1}')"
  printf "\n%10s\t%25s\t%7s\n" "" "Grand Total" "$total"

  popd &>/dev/null
}

action_resquash_now () {
# Args: none
# Returns: none

  [[ $UID != 0 ]] && die "Must be root to perform a resquash"

  info "Voluntary resquash requested"

  local number_of_bins=$(grep -vE "^[ ]*$" "$BINVENTORY" | wc -l)
  if [[ $number_of_bins -le $MIN_BINS ]]; then
    die "Nothing to do. Current backups do not exceed MIN_BINS value."
  else
    create_new_squash $(( $number_of_bins - $MIN_BINS ))
  fi

  exit $?
}

action_restore () {

  [[ -z $1 ]] && die "Missing parameter from restore. Provide a full path to the restore target."
  [[ $UID != 0 ]] && die "Must be root to restore."

  mount_squash || die "Failed to mount seed"

  IFS=$'\n' read -r -d $'\0' -a results < <(find $BINS_DIR/*/$(dirname $1)/$(basename $1) -maxdepth 0 2>/dev/null)
  [[ -e "$SQUASH_MOUNT/$1" ]] && results[0]="$SQUASH_MOUNT/$1"

  [[ ${#results[@]} -eq 0 && ! -e $seedfile ]] && unmount_squash && die "Target not found: '$1'"
  local restore_type=$(stat -c %F ${results[0]})

  declare -a snaps
  [[ -n ${results[0]} ]] && snaps[0]=$(stat -c %Z ${results[0]})

  info "Found $(basename $1) in the following backups:"
  [[ -n ${results[0]} ]] && printf "  0\t%s\n" "$(date --date=@${snaps[0]})"
  for result in "${results[@]:1}"; do
    local bin=$(sed -n "s|$BINS_DIR/\([0-9]*\)$1|\1|p" <<< "$result")
    local bkupdate=$(sed -n "s|^$bin:\([0-9]*\)$|\1|p" $BINVENTORY)
    snaps[$bin]=$bkupdate

    printf "  %d\t%s\n" $bin "$(date --date="@$bkupdate")"
  done

  read -p "Which snapshot to restore from? " reply
  [[ -z $reply || -z ${snaps[$reply]} ]] && die "Invalid bin"

  local bin_list=($(sed -n 's/[\t\ ]*\([0-9]*\):.*/\1/p' "$BINVENTORY"))
  local num=$(grep -nh "^$reply:" $BINVENTORY | cut -d: -f1)

  mountpoint -q $UNION_MOUNT && unmount_union

  [[ $reply != 0 ]] &&
    { mount_union_with_bins ${bin_list[@]:0:$num} || die "Failed to mount union"; } ||
    mount_union_with_bins || die "Failed to mount union"

  local restore_name="$1.$(date --date=@${snaps[$reply]} +%Y%m%d)"

  {
    if [[ -d "$UNION_MOUNT/$1" ]]; then
      rsync -a $UNION_MOUNT/$1/* "$restore_name" &&
      chown $(stat -c %u:%g "$UNION_MOUNT/$1") "$restore_name";
    else
      rsync -a "$UNION_MOUNT/$1" "$restore_name";
    fi
  } && info "Your ${restore_type##regular } has been restored as ${restore_name}"

  unmount_all

}

action_unmount () {
  unmount_all
}

usage () {
    info "SquashFu: Super Awesome Backup Express (Professional Edition)"
    echo "version: $VER"
    cat <<HELP

USAGE
    squashfu <action> [options]

ACTIONS
  -B
    Runs a regular backup, using the config file at /etc/squashfu.

  -C
    Create a new squash by merging old bins. This will still leave you
    with $MIN_BINS backups (as per the MIN_BINS setting in your config).

  -D <bin number>
    Delete an incremental backup. This is done interactively and you will have
    a chance to confirm before any files are deleted.

  -G <path>
    Directly restore a file or directory. This is an interactive operation and
    a list of locations where the target is found will be presented.

  -Q
    Displays the size of the seed, the incrementals, and totals.

  -R <number of bins>
    Rollback specified number of backups and mount union for browsing. The rolled
    back data will be mounted at $UNION_MOUNT.

  -U
    Unmount squash and union. Although SquashFu will always check and unmount as 
    necessary before an operation, this is provided as a safeguard.

OPTIONS
  -c <path_to_config>
    Specify an alternate config file. Options defined will override options in
    the default config. If you specify alternate locations via a supplmenetal config,
    you will need to provide the config for all actions

  -v
    Show debugging output.

HELP
  exit 1
}

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

while getopts :BG:CD:QR:Uc:v opt; do
  case $opt in
    B)
      [[ -n $action ]] && die "only one operation may be used at a time"
      action=backup ;;
    C)
      [[ -n $action ]] && die "only one operation may be used at a time"
      action=resquash_now ;;
    D)
      [[ -n $action ]] && die "only one operation may be used at a time"
      action="remove_bin $OPTARG" ;;
    G)
      [[ -n $action ]] && die "only one operation may be used at a time"
      action="restore $OPTARG" ;;
    Q)
      [[ -n $action ]] && die "only one operation may be used at a time"
      action=report ;;
    R)
      [[ -n $action ]] && die "only one operation may be used at a time"
      action="rollback $OPTARG" ;;
    U)
      [[ -n $action ]] && die "only one operation may be used at a time"
      action=unmount ;;
    c)
      [[ -f $OPTARG ]] && source $OPTARG ;;
    v)
      DEBUG=true ;;
    \:)
      die "Argument missing from -$OPTARG"
      usage ;;
    \?)
      die "Unrecognized option -$OPTARG"
      usage ;;
  esac >&2
done

[[ -z $action ]] && usage

action_$action