#!/bin/bash
# backup-zfs: use zfs send/recv to push/pull snapshots - New does not run twice
prog="$(basename "$0")"

# Function to display usage
usage() {
    cat >&2 <<-EOF
    usage: $prog [-hvq] [-t tag] [-k keep] [-d dateopts] src dest
      use zfs send/recv to push/pull snapshots

      !This Release is default set to send raw. This means you can send encrypted and compressed Datasets 1:1!

      src          the source fs, specified as [host:]pool/path/to/fs
      dest         the destination fs parent, specified as [host:]pool/path/to/fs
                   (the final path component of src will be appended to dest)
      -p           ssh port
      -h           help
      -v           verbose mode
      -C           SSH Compression - obsolete if send_opts="-v -w" are set to raw
      -R           ZFS Re-Init
      -I           Send intermediate Snapshots
      -q           quiet mode
      -t tag       tag to use for naming snapshots (default: backup-zfs)
      -k keep      number of snapshots to keep on src (default: 5)
      -d dateopts  options for date(1) - used to name the snapshots (default: +%F_%T)
      -s           store mode - output snaps from local fs to ssh server
      -r           read mode - read snaps from ssh server to local fs
      -g gpg-id    gpg recipient key id (store mode only)
    EOF
    exit $1
}

# Log function for syslog and stdout (if verbose or tty)
log() {
    logger -t "$prog" -- "$@"
    if ! $quiet && [[ -t 1 ]] || $verbose ; then
        echo "$@" >&2
    fi
}

# Exit with a code & message
die() {
    code="$1"
    shift
    if [[ $code -ne 0 ]] ; then
        verbose=true log "FATAL: $@"
    else
        log "$@"
    fi
    exit $code
}

# Run zfs(1) command either locally or via ssh
ZFS() {
    host="$1"
    shift
    if [[ -n $host ]] ; then
        log "remote ($host): zfs $@"
        ssh $sshcompress -p $port "$host" zfs "$@"
    else
        log "local: zfs $@"
        zfs "$@"
    fi
}

# Create a snapshot on the specified filesystem
create_snapshot() {
    local host=$1
    local fs=$2
    local snapshot_name=$3
    ZFS "$host" snapshot -r "$fs@$snapshot_name" || die $? "zfs snapshot failed"
}

# Send or receive ZFS snapshot between source and destination
ZFS_send_receive() {
    local src=$1
    local dest=$2
    local options=$3
    local action=$4

    log "$action: $src to $dest"
    if [[ $action == "send" ]]; then
        ZFS "$srchost" send $options "$src" | ZFS "$desthost" receive $options -Fue "$dest" || die $? "zfs send failed"
    elif [[ $action == "receive" ]]; then
        ZFS "$srchost" send $options "$src" | ZFS "$desthost" receive $options -Fue "$dest" || die $? "zfs receive failed"
    fi
}

# Encrypt the snapshot using GPG
encrypt_snapshot() {
    local gpg_id=$1
    local snapshot_file=$2
    gpg --trust-model always --encrypt --recipient "$gpg_id" "$snapshot_file" || die $? "GPG encryption failed"
}

# Cleanup old snapshots on the source filesystem
cleanup_old_snapshots() {
    local fs=$1
    local keep=$2
    local tag=$3

    # Get all snapshots for the dataset
    snapshots=$(ZFS "" list -d 1 -t snapshot -H -o name "$fs" | grep -F "@${tag}_" | cut -f2 -d@)
    # Calculate the snapshots to keep
    old_snapshots=$(echo "$snapshots" | tail -n +$((keep+1)))

    # Delete the old snapshots
    for snap in $old_snapshots; do
        ZFS "" destroy -r "$fs@$snap" || die $? "zfs snapshot cleanup failed"
    done
}

# Locking function to prevent multiple instances
lock() {
    local retries=0
    while ! flock -n 9; do
        retries=$((retries + 1))
        if [[ $retries -ge 5 ]]; then
            die 1 "Unable to acquire lock after 5 attempts"
        fi
        sleep 1
    done
}

unlock() {
    flock -u 9
}

# Main processing
(
    flock -n 9 || exit 1

    # Defaults
    tag="$prog"
    dateopts="+%F_%T"
    keep=5
    verbose=false
    quiet=false
    tossh=false
    fromssh=false
    port=22
    sshcompress=""
    reinit=""
    intermediate="-i "

    # Parse options
    while getopts "hvqk:p:t:d:srCIRg:" opt ; do
        case $opt in
            h) usage 0 ;;
            v) verbose=true; send_opts="-v -w"; recv_opts="-v" ;;
            q) quiet=true ;;
            k) keep=$OPTARG ;;
            p) port=$OPTARG ;;
            t) tag=$OPTARG ;;
            d) dateopts=$OPTARG ;;
            s) tossh=true ;;
            r) fromssh=true ;;
            C) sshcompress="-C" ;;
            R) reinit="-R" ;;
            I) intermediate="-I" ;;
            g) gpgid="$OPTARG" ;;
            *) usage 1 ;;
        esac
    done
    shift $((OPTIND-1))

    # Validate mode conflicts
    if $tossh && $fromssh; then
        die 1 "-s and -r are mutually exclusive"
    fi
    if ! $tossh && [[ -n $gpgid ]] ; then
        die 1 "-g can only be used with -s"
    fi

    # Parse src & dest host/fs info
    if [[ $1 =~ :.*: || $2 =~ :.*: ]]; then
        die 1 "invalid fsspec: '$1' or '$2'"
    fi
    if [[ -z $1 || -z $2 ]]; then
        usage 1
    fi
    src="$1"
    dest="$2"

    # SSH mode (store or read)
    if $tossh; then
        log "sending from local zfs filesystem to SSH server"
        create_snapshot "$srchost" "$srcfs" "$snap"

        # Incremental send if applicable
        last_snapshot="$(ssh "$desthost" zfslast)"
        if [[ -z $last_snapshot ]]; then
            die 1 "No valid incremental path from $src to $dest"
        fi

        ZFS_send_receive "$last_snapshot" "$cur" "$send_opts" "send"

        exit
    elif $fromssh; then
        log "receiving from SSH server to local zfs filesystem"
        for file in $(ssh "$srchost" zfsfind | sort); do
            log "receiving $file from $srchost"
            if [[ $file =~ \.gpg$ ]]; then
                ssh "$srchost" zfsget "$file" | gpg | ZFS "$desthost" receive $recv_opts -Fue "$dest" \
                    && ssh "$srchost" rm "$file"
            else
                ssh "$srchost" zfsget "$file" | ZFS "$desthost" receive $recv_opts -Fue "$dest" \
                    && ssh "$srchost" rm "$file"
            fi
        done
        exit
    fi

    # Main snapshot creation and transfer process
    cur="$srcfs@${tag}_$date"
    create_snapshot "$srchost" "$srcfs" "$cur"

    last_snapshot="$(ZFS "$desthost" list -d 1 -t snapshot -H -S creation -o name "$destfs" | head -n 1)"
    if [[ -z $last_snapshot ]]; then
        ZFS_send_receive "$cur" "$destfs" "$send_opts" "send"
    else
        ZFS_send_receive "$last_snapshot" "$cur" "$send_opts" "send"
    fi

    # Clean up old snapshots on source
    cleanup_old_snapshots "$srcfs" "$keep" "$tag"
) 9>/var/lock/bashclub-zfs.lock
