#!/usr/bin/env sh

# requires zfs, pv, ssh

nodes_file="nodes.txt"
### nodes.txt
# <origin host>   <user> <ssh port> <timezone> <compressed flag>
# <mirror 1 host> <user> <ssh port> <zfs root> <send speed>
# <mirror 2 host> <user> <ssh port> <zfs root> <send speed>
# ...

datasets_file="datasets.txt"
### datasets.txt
# <dataset 1> [options]
# <dataset 2> [options]
# <dataset 3> [options]
# ...
# recognized options: nosend nosnap(shot) remotezfs=<remote dset name>

if [ ! -f $nodes_file ]; then echo "No nodes specified."; exit 2; fi
if [ ! -f $datasets_file ]; then echo "No datasets specified."; exit 3; fi

ZFS_SEND_OPTS="-pRLe"
ZFS_SNAPSHOT_OPTS="-r"
ZFS_RECV_OPTS="-vsu"
SSH_SEND_OPTS="-Ti ~/.ssh/id_rsa"
DEFAULT_SPEED="20M"

read LOCAL_HOST LOCAL_USER LOCAL_PORT LOCAL_TZ SEND_COMPRESSED < $nodes_file

ORIGIN_PERMS="send,snapshot,hold"
MIRROR_PERMS="compression,receive,create,mount,mountpoint,readonly,jailed"

SEND_DATASETS=""
SNAPSHOT_DATASETS=""

debug()
{
    if [ "$ZB_VERBOSE" = "TRUE" ]; then
        echo $@
    fi
}

run_send_recv()
{
    ssh_cmd="ssh $SSH_SEND_OPTS -o port=$port $user@$addr"
    if [ $ZB_TEST_MODE = "TRUE" ]; then
        echo "$send_cmd | $pv_cmd | $ssh_cmd $recv_cmd"
    else
        $send_cmd | $pv_cmd | $ssh_cmd $recv_cmd
    fi
}

run_snapshot()
{
    if [ $ZB_TEST_MODE = "TRUE" ]; then
        echo zfs snapshot $ZFS_SNAPSHOT_OPTS $dset@$snapid
    else    
        zfs snapshot $ZFS_SNAPSHOT_OPTS $dset@$snapid
    fi
}

run_allow()
{
    if [ $ZB_TEST_MODE = "TRUE" ]; then
        echo zfs allow -u $LOCAL_USER $ORIGIN_PERMS $dset
    else
        zfs allow -u $LOCAL_USER $ORIGIN_PERMS $dset
    fi
}

run_remote()
{
    ssh_cmd="ssh $SSH_SEND_OPTS -o port=$port $user@$addr"
    if [ $ZB_TEST_MODE = "TRUE" ]; then
        echo "\`$ssh_cmd $remote_cmd\`"
    else
        $ssh_cmd $remote_cmd
    fi
}

set_compression() {
    case $SEND_COMPRESSED in
        [tTyY]*) ZFS_SEND_OPTS="$ZFS_SEND_OPTS -c"
        debug "Enabled send as compressed datasets." ;;
    esac
}

populate_datasets()
{
    while read dset dataset_options; do
        case $dset in
            \#*) continue;;
        esac

        snapshot_dataset="yes"
        send_dataset="yes"
        for option in $dataset_options; do
            case $option in 
                "nosnap"*    )
                    snapshot_dataset=""
                    ;;
                "nosend"*    )
                    send_dataset=""
                    ;;
                "remotezfs="*)
                    remotezfs=`echo $option|cut -d'=' -f2`
                    dset="$dset:$remotezfs"
                    ;;
            esac
        done

        if [ $send_dataset ]; then
            SEND_DATASETS="$SEND_DATASETS $dset"
            debug "Enabled send for $dset."
        fi
        if [ $snapshot_dataset ]; then
            SNAPSHOT_DATASETS="$SNAPSHOT_DATASETS $dset"
            debug "Enabled snapshot for $dset."
        fi
    done < datasets.txt
}

list_datasets()
{
    for dset in $SEND_DATASETS; do
        dset=`echo $dset|cut -d: -f1`
        zfs list -H $dset
    done
}

list_remote_datasets() {
    while read addr user port zroot speed; do
        echo "=== $user@$addr $zroot ==="
        if [ $zroot ]; then
            case $zroot in 
                *"/") ;;
                *) zroot="$zroot/" ;;
            esac
        fi
        remote_cmd="
            sh -c 'for dset in $SEND_DATASETS; do
                dset=`echo \\$dset|cut -d: -f2`
                zfs list -H $zroot\$dset;
            done'
        "
        run_remote
    done < sendnodes.tmp
}

get_latest_snapshot()
{
    zfs list -t snapshot $dset | tail -n1 | cut -d'@' -f2 | cut -d' ' -f 1
}

get_latest_remote_snapshot()
{
    if [ $zroot ]; then
        case $zroot in 
            *"/") ;;
            *) zroot="$zroot/" ;;
        esac
    fi
    list_cmd="zfs list -t snapshot $zroot$dset"
    remote_cmd="$list_cmd|tail -n1|cut -d'@' -f2|cut -d' ' -f 1"
    run_remote
        
}

zfs_allow_origin()
{
    for dset in $SEND_DATASETS; do
        dset=`echo $dset|cut -d: -f1`
        run_allow
    done
}

zfs_allow_mirrors()
{
    while read addr user port zroot speed; do
        echo "=== $user@$addr $zroot ==="
        if [ $zroot ]; then
            case $zroot in 
                *"/") ;;
                *) zroot="$zroot/" ;;
            esac
        fi
        remote_cmd="
            for dset in $SEND_DATASETS; do
                dset=`echo \\$dset|cut -d: -f2`
                zfs allow -u $user $MIRROR_PERMS $zroot\$dset
            done
        "
        run_remote
    done < sendnodes.tmp
}

create_snapshots()
{
    snapid="`date -I`-`date +%s`-$LOCAL_TZ"
    for dset in $SNAPSHOT_DATASETS; do
        run_snapshot
        debug Created snapshot $dset@$snapid.
    done
}

send_latest()
{
    while read addr user port zroot speed; do
        echo "=== $user@$addr $zroot ==="
        if [ $zroot ]; then
            case $zroot in 
                *"/") ;;
                *) zroot="$zroot/" ;;
            esac
        fi
        for dset in $SEND_DATASETS; do
            rset=`echo $dset|cut -d: -f2`
            dset=`echo $dset|cut -d: -f1`
            debug "Processing $dset (remote $rset)."
            snap1=`get_latest_snapshot $dset`
            send_cmd="zfs send $ZFS_SEND_OPTS $dset@$snap1"
            pv_cmd="pv -qL ${speed:-$DEFAULT_SPEED}"
            recv_cmd="zfs recv $ZFS_RECV_OPTS $zroot$rset"
            run_send_recv
        done
    done < sendnodes.tmp
}

send_increment()
{
    while read addr user port zroot speed; do
        echo "=== $user@$addr $zroot ==="
        if [ $zroot ]; then
            case $zroot in 
                *"/") ;;
                *) zroot="$zroot/" ;;
            esac
        fi
        for dset in $SEND_DATASETS; do
            rset=`echo $dset|cut -d: -f2`
            dset=`echo $dset|cut -d: -f1`
            snap1=`get_latest_snapshot`
            snap0=`get_latest_remote_snapshot`
            send_cmd="zfs send $ZFS_SEND_OPTS -i $dset@$snap0 $dset@$snap1"
            pv_cmd="pv -qL ${speed:-$DEFAULT_SPEED}"
            recv_cmd="zfs recv $ZFS_RECV_OPTS $zroot$rset"
            run_send_recv
        done
    done < sendnodes.tmp
}


################################### MAIN ####################################

### Options
set -- `getopt "tvl:" "$@"` || {
    echo "Usage: `basename $0` [-tvl] [command]" 1>&2
    exit 1
}
ZB_TEST_MODE=FALSE
while :; do
    case "$1" in
        -t)
            ZB_TEST_MODE=TRUE
            echo Test mode.
        ;;
        -v)
            ZB_VERBOSE=TRUE
            ZFS_SEND_OPTS="$ZFS_SEND_OPTS -v"
            # Current on by default.
            # ZFS_RECV_OPTS="$ZFS_RECV_OPTS -v"
            echo Verbose mode.
        ;;
        --) shift; break ;;
    esac
    shift
done

tail -n +2 $nodes_file|grep -v '^#' > sendnodes.tmp

set_compression
populate_datasets

### Command
case "$1" in
    list              ) list_datasets ;;
    list-remote       ) list_remote_datasets ;;
    snapshot          ) create_snapshots ;;
    configure-mirrors ) zfs_allow_mirrors ;;
    configure-origin  ) zfs_allow_origin ;;
    mirror-reset      ) send_latest ;;
    mirror-increment  ) send_increment ;;
    #mirror-prune      ) send_increment_with_F ;;
esac

rm sendnodes.tmp