#!/usr/bin/bash
#
# Copyright(C) 2020 Benoit DOLEZ @ ZENETYS
# This script is licensed under MIT License (http://opensource.org/licenses/MIT)
#

export LC_ALL=C
shopt -s nullglob
set -o pipefail

ARGV0=${BASH_SOURCE[0]}
PROGNAME=${0##*/}
HOSTNAME=$(< /proc/sys/kernel/hostname)
HOSTNAME=${HOSTNAME%%.*}
DRY_RUN=1
BECOMEROOT=0
LOCK_FILE="/dev/shm/$PROGNAME.lock"
LOCK_TAKEN=
declare -r ZSYNC_TEMPDIR="${TMPDIR:-/tmp}/$PROGNAME.$$.$RANDOM"

CONFIG=${ZSYNC_CONFIG:-$HOME/.zsync.cfg}
[[ -r $CONFIG ]] || CONFIG=/etc/ztools/zsync.cfg
TIMEOUT=10
PEERS=()
INCLUDES=()
EXCLUDES=()
USE_RSYNC_UPDATE= # set USE_RSYNC_UPDATE=1 in config to use rsync --update
USE_RSYNC_XATTRS= # set USE_RSYNC_XATTRS=1 in config to use rsync --xattrs
SUDO_REMOTE=""

if tty -s >/dev/null 2>&1; then
    declare -A C=(
        [info]=$'\x1b[1m'
        [warning]=$'\x1b[7;33m'
        [error]=$'\x1b[7;91m'
        [fatal]=$'\x1b[7;1;91m'
        [blue]=$'\x1b[7;34m'
        [reset]=$'\x1b[m'
    )
    HAS_TTY=1
else
    declare -A C=()
    HAS_TTY=
fi

function exit_usage() {
    local status=${1:-0}
    [[ "$status" != "0" ]] && exec >&2

    echo "\
Usage: $PROGNAME [OPTION]... [--] [FILE]...
Rsync wrapper for cluster/HA platforms

Options:
    -f, [--]force       Disable dry-run mode
    -t, [--]check       Check mode, implies dry-run
    -d, [--]diff        View diff on text files
    -c, --config        Path to configuration file
    -H, --host          Set target host
    -i, --include       Add include rule
    -e, --exclude       Add exclude rule
    -q, --quiet         Disable stdout, repeat for stderr
    -h, --help          Display this help

Current configuration file: $CONFIG

Use of options -i and/or -e ignores the include/exclude rules
defined in the configuration file. Similarly, use of option -H
ignores the peers defined in the configuration file.

These two forms are equivalent:
    $PROGNAME -i INC1 -i INC2 -e EXC1
    $PROGNAME -e EXC1 INC1 INC2

Check mode is ment to be used from scripts to test if peers are
in sync. In check mode, status codes are as follows:
    0: peers in sync
    1: peers out of sync
    2: error, sync status is undefined or fatal error

Flag -q, --quiet is normally ment to be used with check mode.
"
    exit "$status"
}

function info() {
    local i;
    local d=$(date +%Y-%m-%dT%H:%M:%S.%3N%z);
    local l=info
    [[ ${1:0:2} == -- ]] && l=${1:2} && shift
    for i in "$@"; do
        printf '%s %s %s%s%s %s: %s\n' "$d" "$HOSTNAME" "${C[$l]}" "$l" "${C[reset]}" "$PROGNAME" "${DRY_RUN:+DRY-RUN ** }$i";
    done
}

function fatal() {
    info --${FUNCNAME[0]} "$@"
    exit 2
}

function error() {
    info --${FUNCNAME[0]} "$@"
}

function warning() {
    info --${FUNCNAME[0]} "$@"
}

function take_lock_or_fatal() {
    if ! ln -sT "/proc/$$" "$LOCK_FILE" 2>/dev/null; then
        [[ -e "$LOCK_FILE" ]] || release_lock
        if ! ln -sT "/proc/$$" "$LOCK_FILE" 2>/dev/null; then
            fatal "Cannot take lock $LOCK_FILE"
        fi
    fi
    LOCK_TAKEN=1
}

function release_lock() {
    rm -f "$LOCK_FILE" || return 1
    LOCK_TAKEN=
}

function need_tempdir() {
    mkdir -p "$ZSYNC_TEMPDIR" || return 1
}

function on_exit() {
    if [[ -n $ZSYNC_TEMPDIR && -d $ZSYNC_TEMPDIR/. ]]; then
        rm -rf "$ZSYNC_TEMPDIR"
    fi
    if [[ -n $LOCK_TAKEN ]]; then
        release_lock
    fi
    echo -n "${C[reset]}"
}

function confirm() {
    local user_input
    echo -n "$1"
    read -r user_input
    if [[ "$user_input" != [Yy] ]]; then
        info "Sync aborted at user request"
        return 1
    fi
}

# $1: PEERS entry
function peer_is_myself() {
    local short_peer_host=${1#*@}; short_peer_host=${short_peer_host%%.*}
    [[ ${short_peer_host,,} == ${HOSTNAME,,} ]] && return 0
    return 1
}

# $1: rsync diff list file
# $2: peer name
function show_diff() {
    local difflist=$1 ; shift
    local peer=$1 ; shift
    local hash action mtime owner file
    local -a ascii_files=( )

    local diffcolortest=$(diff --color 2>&1)
    if (( $? == 127 )); then
        error 'diff command not found'
        return 1
    elif [[ $diffcolortest != *'unrecognized option'* ]]; then
        local diffcmd=( diff --color=auto )
    else
        local diffcmd=( diff )
    fi

    while read hash acode action mtime perm owner file ; do
        printf "# %s %s %s %s %s %s\n" \
               "$acode" "$action" \
               "$mtime" "$perm" "$owner" "$file"
        if [[ ${acode:1:2} == "fc"     ]] &&
           (( ${#ascii_files[@]} <= 10 )) &&
           [[ $(file $file) =~ ASCII|text ]]; then
            ascii_files+=( "$file" )
            (( ${#ascii_files[@]} > 10 )) &&
                warning "Too many text files for diff, display the first 10"
        fi
    done < "$difflist"

    if (( ${#ascii_files[@]} > 0 )); then
        # echo "ASCII : ${ascii_files[*]}" >&2
        if ! need_tempdir; then
            error "can't create temporary directory"
            return 1
        fi
        mkdir -p "$ZSYNC_TEMPDIR/old"
        ln -s / "$ZSYNC_TEMPDIR/new"
        (
            IFS=$'\t'
            function do_remote() {
                tar zc -C / "${@#/}"
            }
            declare -f do_remote
            echo "do_remote$IFS${ascii_files[*]}"
        ) |
            "${SSHCMD[@]}" "$peer" "[[ -x $0 ]] && CMD='$0 --do-remote' || CMD='bash'; $SUDO_REMOTE \$CMD" |
            tar zx -C "$ZSYNC_TEMPDIR/old"
        (
            cd "$ZSYNC_TEMPDIR"
            find old -type f -printf "%P\0" |
                xargs -0 -I{} "${diffcmd[@]}" --unidirectional-new-file -u old/{} new/{}
        )
        rm -rf "$ZSYNC_TEMPDIR"/{old,new}
    fi
}

function do_remote() {
    local IFS=$'\t'
    local -a FILES=( )
    while read CMD REST; do
       [[ $CMD == do_remote ]] && FILES=( $REST )
    done
    if [[ ${#FILES[@]} -gt 0 ]]; then
        tar zc -C / "${FILES[@]}"
    fi
    return $?
}

function do_remote_rsync() {
    [[ $1 == --server ]] || return 1
    rsync "$@"
    return $?
}

function do_rsync() {
    # retval: This function status code, non-zero on error.
    # CHECK_OUT_OF_SYNC: Only in check mode. Track peers that are out of sync.
    #   This variable is populated only in check mode.
    CHECK_OUT_OF_SYNC=()
    local i retval=0 peer peer_retval
    local rsync_opts=(
        --archive # -rlptgoD
        --acls
        ${USE_RSYNC_UPDATE:+--update}
        ${USE_RSYNC_XATTRS:+--xattrs}
        --one-file-system
        --compress
        --checksum
        --relative
        --omit-dir-times
        --delete
        --rsync-path "[[ -x $ARGV0 ]] && CMD='${ARGV0@Q} --do-rsync' || CMD='rsync' ; $SUDO_REMOTE \$CMD"
        -e "${SSHCMD[*]}"
    )

    local dry_run_format=' # %i %o %M %B %U:%G  /%f'
    local sync_format=' > %i %o /%f'
    if [[ -n $HAS_TTY || -n $DRY_RUN ]] && ! need_tempdir; then
        error "Can't create temporary directory"
        return 2
    fi

    for i in "${EXCLUDES[@]}"; do
        rsync_opts+=( --exclude "$i" )
    done

    for peer in "${PEERS[@]}"; do
        # bypass myself
        peer_is_myself "$peer" && continue

        # list or diff changes
        if [[ -n $HAS_TTY || -n $DRY_RUN ]]; then
            info "Check sync of ${C[info]}${peer}${C[reset]}"
            rsync "${rsync_opts[@]}" --dry-run --out-format "$dry_run_format" "${INCLUDES[@]}" "$peer:/" > "$ZSYNC_TEMPDIR/list"
            peer_retval=$?

            if (( peer_retval != 0 )); then
                error "Rsync to $peer returned error code $peer_retval"
                if (( peer_retval == 255 )); then
                    error "Try: ssh-copy-id${SSH_USER:+ -l "$SSH_USER"}${SSH_IDENTITY:+ -i "$SSH_IDENTITY"} -o StrictHostKeyChecking=ask $peer"
                fi
                retval=1
                continue
            fi

            # bypass if diff empty
            [[ -s $ZSYNC_TEMPDIR/list ]] || continue

            if [[ -n $DIFF ]]; then
                show_diff "$ZSYNC_TEMPDIR/list" "$peer"
            else
                cat "$ZSYNC_TEMPDIR/list"
            fi

            # register peer as out of sync in check mode
            [[ -n $CHECK ]] && CHECK_OUT_OF_SYNC+=( "$peer" )
            # bypass in dry-run mode
            [[ -n $DRY_RUN ]] && continue
            # otherwise there is a tty and user must confirm to apply changes
            confirm "Please confirm to ${C[error]}sync${C[reset]} files? [N/y] " || continue
        fi

        info "Rsync to $peer"
        rsync "${rsync_opts[@]}" --out-format "$sync_format" "${INCLUDES[@]}" "$peer:/"
        peer_retval=$?

        if (( peer_retval == 0 )); then
            info "Done rsync to $peer"
        else
            error "Rsync to $peer returned error code $peer_retval"
            retval=1
        fi
    done
    rm -f "$ZSYNC_TEMPDIR/list"
    return "$retval"
}

peers=()
includes=()
excludes=()
use_rsync_update=
ARGV=( "$@" )
XDEBUG=
DIFF=
CHECK=
QUIET=0

while (( $# > 0 )); do
    case "$1" in
        -f|--force|push) DRY_RUN= ;;
        -t|--check|check) CHECK=1 ;;
        -d|--diff|diff) DIFF=1 ;;
        -c|--config) CONFIG=$2; shift ;;
        -R|--become-root) BECOMEROOT=1 ;;
        -H|--host) peers+=( "$2" ); shift ;;
        -i|--include) includes+=( "$2" ); shift ;;
        -e|--exclude) excludes+=( "$2" ); shift ;;
        --use-rsync-update) use_rsync_update=1 ;;
        --do-remote) do_remote; exit $? ;;
        --do-rsync)  do_remote_rsync "${@:2}"; exit $? ;;
        --x-debug) XDEBUG=1 ;;
        -q|--quiet) (( QUIET++ )) ;;
        -h|--help) exit_usage ;;
        --) shift; break ;;
        -*) exit_usage 1 ;;
        *)  includes+=( "$1" ) ;;
    esac
    shift
done

includes+=( "$@" )

[[ -n $XDEBUG ]] && set -x
(( QUIET >= 1 )) && exec 1>/dev/null
(( QUIET >= 2 )) && exec 2>/dev/null

trap on_exit EXIT

if [[ -n $CONFIG && -r $CONFIG ]]; then
    if ! source "$CONFIG"; then
        fatal "Cannot load configuration file '$CONFIG'"
    fi
fi

# Allow overriding rsync --update from command line
[[ -n $use_rsync_update ]] && USE_RSYNC_UPDATE=$use_rsync_update

# This script is ment to be run as root.
if (( $BECOMEROOT == 1 && $UID != 0 )); then
    sudo_opts=( -u root )
    [[ -t 0 ]] || sudo_opts=( -n )
    exec sudo -u root "${sudo_opts[@]}" -- "$0" "${ARGV[@]}"
elif (( $BECOMEROOT == 1 )); then
    SUDO_REMOTE='[ "$UID" != 0 ] && SUDO="sudo -n -- "; $SUDO'
fi

if (( ${#peers[@]} > 0 )); then
    PEERS=( "${peers[@]}" )
fi

[[ -z $PEERS ]] && fatal "No peers to sync with"

[[ $SSH_IDENTITY && ! -r $SSH_IDENTITY ]] &&
  fatal "Cannot read SSH identity '$SSH_IDENTITY'"

SSHCMD=(
    ssh -T
    ${SSH_USER:+-l "$SSH_USER"}
    ${SSH_IDENTITY:+-o IdentityFile="$SSH_IDENTITY"}
    ${SSH_LOGLEVEL:+-o LogLevel="$SSH_LOGLEVEL"}
    -o ConnectTimeout=$TIMEOUT
    -o IdentitiesOnly=yes
    -o BatchMode=no
    -o ClearAllForwardings=yes
    -o CheckHostIP=yes
    -o StrictHostKeyChecking=ask
)

(( ${#includes[@]} > 0 )) && INCLUDES=( "${includes[@]}" )
(( ${#excludes[@]} > 0 )) && EXCLUDES=( "${excludes[@]}" )

[[ -z $INCLUDES ]] && fatal "Noting to include for sync"

for i in "${INCLUDES[@]}"; do
    if [[ ${i:0:1} != / ]]; then
        fatal "Expected absolute path for include '$i'"
    fi
done

[[ -n $CHECK ]] && DRY_RUN=1

[[ -z $DRY_RUN ]] && take_lock_or_fatal

do_rsync
retval=$?

if [[ -n $CHECK ]]; then
    case "$retval,${#CHECK_OUT_OF_SYNC[@]}" in
        0,0) exit 0 ;;
        0,*) exit 1 ;;
        *) exit 2 ;; # same as fatal
    esac
fi
exit "$retval"
