#!/bin/bash
#
# bashclub zfs replication script
# Author: (C) 2024 Thorsten Spille <thorsten@spille-edv.de>

# version is taken from the installed package (set by the release tag)
__version__=v$(dpkg-query -W -f='${Version}' bashclub-zsync 2>/dev/null || echo dev)

# also search next to this script (manual installation, e.g. on a TrueNAS dataset)
PATH="$(dirname "$(readlink -f "$0")"):/bin:/sbin:/usr/bin:/usr/sbin:/usr/local/bin:/usr/local/sbin"

prog="$(basename $0)"
checkzfs=$(which checkzfs)
rsm=$(which remote-snapshot-manager)

rc=0
debug=

#### default config file, can be changed with parameter -c
conf=/etc/bashclub/zsync.conf

#### default values in config file
# replication target path on local machine 
target=pool/dataset

# ssh address of remote machine
source=user@host

# ssh port of remote machine
sshport=22

# zfs user parameter to identify filesystems/volumes to replicate
tag=bashclub:zsync

# pipe separated list of snapshot name filters
snapshot_filter="hourly|daily|weekly|monthly"

# minimum count of snapshots per filter to keep
min_keep=3

# number of zfs snapshots to keep on source (0 or 1 = snapshot function disabled)
zfs_auto_snapshot_keep=0

# make snapshot via zfs-auto-snapshot before replication
zfs_auto_snapshot_label="backup"

# select snapshot engine: "zas" or "internal"
# keep "zas" here for existing configs without this setting; new configs get "internal"
zfs_auto_snapshot_engine="zas"

# disable checkzfs with value > 0
checkzfs_disabled=0

# do not connect to source host (if replication is local)
checkzfs_local=0

# set checkzfs parameter "--prefix"
checkzfs_prefix=zsync

# set checkzfs maximum age of last snapshot in minutes (comma separated => warn,crit)
checkzfs_max_age=1500,6000

# set checkzfs maximum count of snapshots per dataset (comma separated => warn,crit)
checkzfs_max_snapshot_count=150,165

# set where to move the checkzfs output (0 = local machine, 1 = source machine)
checkzfs_spool=0

# set maxmimum age of checkzfs data in seconds (hourly = 3900, daily = 87000)
checkzfs_spool_maxage=87000

usage() {
	cat >&2 << EOF
$prog Version $__version__
-------------------------------------------------------------------------------
usage: $prog [-h] [-d] [-c CONFIG]
    creates a mirrored replication of configured zfs filesystems / volumes
    -c CONFIG    Configuration file for this script
    -d           Debug mode
-------------------------------------------------------------------------------
    (C) 2024     by Spille IT Solutions for bashclub (github.com/bashclub)
                 Author: Thorsten Spille <thorsten@spille-edv.de>
-------------------------------------------------------------------------------
EOF
	exit $1
}

args=("$@")
while getopts "hdc:" opt; do
  case $opt in
    h) usage 0 ;;
    c) conf=$OPTARG ;;
    d) debug=-v ;;
    *) usage 1 ;;
  esac
done
shift $((OPTIND-1))

function log() {
    echo -e "$(date +'%b %d %T') $1"
}

function debug() {
    if [[ $debug == "-v" ]]; then log "[DEBUG] $1"; fi
}

# run a command on the source host
function remote() {
    ssh $controlmaster $controlpath $controlpersist $sshcipher $zsync_sshport $source "$@"
}

#### one run per config at a time: the script runs again under flock, a second run is skipped
# flock -o: the script itself does not hold the lock, so the background ssh master
# connection (ControlPersist) does not keep the lock after the run has finished
# only flock from util-linux (Linux): the options -o and -E are not available everywhere
if [[ -z $ZSYNC_LOCK ]]; then
    if flock --version 2>/dev/null | grep -q util-linux; then
        lockdir=/run/lock
        [ -d $lockdir ] || lockdir=/tmp
        export ZSYNC_LOCK=$lockdir/bashclub-zsync$(readlink -m "$conf" | tr '/' '_').lock
        flock -n -o -E 75 "$ZSYNC_LOCK" "$BASH" "$0" "${args[@]}"
        lock_rc=$?
        if [ $lock_rc -eq 75 ]; then
            log "[ERROR] Another run with configuration $conf is still active (pid $(cat "$ZSYNC_LOCK" 2>/dev/null)). Skipping this run."
            exit 1
        fi
        exit $lock_rc
    fi
    log "[WARN] flock (util-linux) not found, running without protection against concurrent runs."
else
    # pid of the active run for the message above
    echo $$ > "$ZSYNC_LOCK"
fi

#### log file without logrotate (e.g. TrueNAS: cron job with ">> /mnt/<pool>/.../bashclub-zsync.log")
# if the output is appended to a file whose directory no logrotate config mentions,
# the script rotates it at 10 MiB and keeps 4 compressed old logs
logfile=$(readlink -f /proc/$$/fd/1 2>/dev/null)
if [[ -f $logfile ]] && [ $(stat -c %s "$logfile") -gt 10485760 ] \
    && ! cat /etc/logrotate.conf /etc/logrotate.d/* 2>/dev/null | grep -qF "$(dirname "$logfile")/"; then
    errfile=$(readlink -f /proc/$$/fd/2 2>/dev/null)
    for i in 3 2 1; do
        if [ -f "$logfile.$i.gz" ]; then mv -f "$logfile.$i.gz" "$logfile.$((i+1)).gz"; fi
    done
    mv -f "$logfile" "$logfile.1"
    exec >> "$logfile"
    if [[ $errfile == "$logfile" ]]; then exec 2>&1; fi
    gzip -f "$logfile.1"
    log "[INFO] Log file rotated, previous log: $logfile.1.gz"
fi

# load config file
if [ -f $conf ]; then
    log "[INFO] Reading configuration $conf"
    source $conf
else
    mkdir -p $(dirname $conf)
    cat << EOF > $conf
target=$target
source=$source
sshport=$sshport
tag=$tag
snapshot_filter="$snapshot_filter"
min_keep=$min_keep
zfs_auto_snapshot_keep=$zfs_auto_snapshot_keep
zfs_auto_snapshot_label=$zfs_auto_snapshot_label
zfs_auto_snapshot_engine=internal
checkzfs_disabled=$checkzfs_disabled
checkzfs_local=$checkzfs_local
checkzfs_prefix=$checkzfs_prefix
checkzfs_max_age=$checkzfs_max_age
checkzfs_max_snapshot_count=$checkzfs_max_snapshot_count
checkzfs_spool=$checkzfs_spool
checkzfs_spool_maxage=$checkzfs_spool_maxage
EOF
    log "[INFO] Initial config file created. Please adjust and restart script. Exiting..."
    usage 0
fi

# checked after loading the config, so checkzfs_disabled from the config is respected
if [[ -n $checkzfs ]]; then
    if [ $checkzfs_disabled -eq 0 ]; then
        log "[INFO] installed checkzfs found: $checkzfs"
    else
        log "[INFO] checkzfs is disabled by configuration."
    fi
else
    checkzfs_disabled=1
    log "[WARN] checkzfs is disabled. for monitoring please install:\n\thttps://code.bashclub.org/bashclub/checkzfs\n"
fi

if [[ -z $source ]]; then
    log "[ERROR] source is empty. For local replication use ssh to localhost, e.g. source=root@localhost"
    exit 1
fi
zsync_sshport=-p$sshport

config_log="[INFO] Configuration:"
for setting in target source sshport tag snapshot_filter min_keep zfs_auto_snapshot_keep zfs_auto_snapshot_label \
               zfs_auto_snapshot_engine checkzfs_disabled checkzfs_local checkzfs_prefix checkzfs_max_age \
               checkzfs_max_snapshot_count checkzfs_spool checkzfs_spool_maxage; do
    config_log+="\n\t$setting=${!setting}"
done
log "$config_log\n"

#### ssh connection: shared master connection, aes cipher if both hosts support it
controlmaster=-oControlMaster=auto
controlpath=-oControlPath=~/.ssh/bashclub-zsync-%r@%h-%p
controlpersist=-oControlPersist=15

local_os_id=$(grep -E "^ID=" /etc/os-release | cut -d'=' -f2)
remote_os_id=$(remote "grep -E \"^ID=\" /etc/os-release | cut -d'=' -f2")

if [[ $local_os_id == "freebsd" ]]; then
    local_aes=$(grep -o AES /var/run/dmesg.boot | uniq)
else
    local_aes=$(grep -m1 -o aes /proc/cpuinfo | uniq)
fi

if [[ $remote_os_id == "freebsd" ]]; then
    remote_aes=$(remote "grep -o AES /var/run/dmesg.boot | uniq")
else
    remote_aes=$(remote "grep -m1 -o aes /proc/cpuinfo | uniq")
fi

if [[ $local_aes == "aes" ]] && [[ $remote_aes == "aes" ]]; then
    debug "Switching to cipher aes256-gcm@openssh.com"
    sshcipher=-caes256-gcm@openssh.com
else
    debug "Using default cipher chacha20-poly1305@openssh.com"
    sshcipher=-cchacha20-poly1305@openssh.com
fi

#### datasets to replicate: zfs user property $tag on the source
# all     = this dataset (if set locally or inherited, not received)
# subvols = only child datasets (the dataset with the local setting itself is skipped)
include_list="[INFO] Included datasets:\n"
exclude_list="[INFO] Excluded datasets:\n"
checkzfs_headline="[INFO] checkzfs filter:"
checkzfs_filter=
IFS=$'\n'
for zvol in $(remote "zfs get -H -o name,value,source -t filesystem,volume $tag"); do
    # property source: local, received, inherited from <dataset> or default
    IFS=$'\t' read -r name value property_source <<< "$zvol"
    if [[ $value == "subvols" ]]; then
        if [[ $property_source != "local" ]] && [[ $property_source != "received" ]]; then
            include_list="${include_list}\t${name}\n"
            syncvols+=("$name")
        elif [[ $property_source == "local" ]]; then
            checkzfs_filter="#$name/|$checkzfs_filter"
        fi
    elif [[ $value == "all" ]]; then
        if [[ $property_source != "received" ]]; then
            include_list="${include_list}\t${name}\n"
            syncvols+=("$name")
        fi
        if [[ $property_source == "local" ]]; then
            checkzfs_filter="#$name|$checkzfs_filter"
        fi
    else
        exclude_list="${exclude_list}\t${name}\n"
    fi
done
checkzfs_filter=${checkzfs_filter::-1}
declare -A fstypes
while IFS=$'\t' read -r dataset fstype; do fstypes[$dataset]=$fstype; done <<< "$(remote "zfs list -H -o name,type -t filesystem,volume")"
log "$include_list"
log "$exclude_list"
log "$checkzfs_headline\ncheckzfs_filter=$checkzfs_filter\n"

#### target dataset, never snapshotted by zfs-auto-snapshot on this host
if ! zfs list $target > /dev/null 2>&1 ; then
    debug "$target does not exist. Creating..."
    zfs create -o canmount=noauto -o com.sun:auto-snapshot=false $target
else
    debug "$target exists, check auto-snapshot..."
    if [[ $(zfs get -H -o value,source com.sun:auto-snapshot $target) != "false	local" ]]; then
        zfs set com.sun:auto-snapshot=false $target
    fi
fi

#### snapshots on the source before replication (zas: once for all datasets, internal: per dataset below)
if [ $zfs_auto_snapshot_keep -gt 1 ]; then
    if [[ $zfs_auto_snapshot_engine == "zas" ]]; then
        debug "Running zfs-auto-snapshot"
        remote "which zfs-auto-snapshot > /dev/null || exit 0 ; zfs-auto-snapshot --quiet --syslog --label=$zfs_auto_snapshot_label --keep=$zfs_auto_snapshot_keep //"
    fi

    if [[ $snapshot_filter == "" ]]; then
        snapshot_filter="$zfs_auto_snapshot_label"
    elif [[ "|$snapshot_filter|" != *"|$zfs_auto_snapshot_label|"* ]]; then
        # only if not already in the filter, otherwise each interval is checked twice
        snapshot_filter="$snapshot_filter|$zfs_auto_snapshot_label"
    fi
fi

#### replication per dataset
for name in "${syncvols[@]}"; do
    log "[INFO] Replicate $name"
    fstype=${fstypes[$name]}
    if [[ $fstype == "filesystem" ]]; then
        mp=-xmountpoint
        cm=-ocanmount=noauto
    else
        mp=
        cm=
    fi

    if [ $zfs_auto_snapshot_keep -gt 1 ]; then
        if [[ $zfs_auto_snapshot_engine == "internal" ]]; then
            debug "Running remote-snapshot-manager"
            $rsm $debug -d $name -k $zfs_auto_snapshot_keep -l $zfs_auto_snapshot_label -h $source -p $sshport
        fi
    fi

    # source snapshots (oldest first), queried directly before the replication of this dataset
    source_snaps=$(remote "zfs list -H -o name,guid -s creation -t snapshot $name 2>/dev/null")
    if cut -f1 <<< "$source_snaps" | grep -qE "@.*($snapshot_filter)"; then
        IFS=$' '
        # initial replication: create the path on the target, then send the oldest matching snapshot
        if ! zfs list -H $target/$name > /dev/null 2>&1 ; then
            debug "$target/$name does not exist"
            prefix=""
            for part in $(echo $target/$(echo $name | cut -d'/' -f1) | sed "s/\// /g"); do
                if [ $(zfs list $prefix$part > /dev/null 2>&1 ; echo $?) -gt 0 ]; then
                    debug "$prefix$part does not exist"
                    log "[INFO] Creating $prefix$part"
                    zfs create -o canmount=noauto $autosnap -p $prefix$part
                fi
                prefix="$prefix$part/"
            done

            debug "$name - Start initial replication"
            IFS=$'\n'
            snap=$(remote "zfs list -H -t snapshot -o name -S creation $name | grep -E \"@.*($snapshot_filter)\" | tail -1")
            if [[ -n $snap ]]; then
                log "[INFO] Start initial replication: $snap => $target/$(echo $name | cut -d'/' -f1)"
                if remote "zfs send -w -p $debug $snap" | zfs receive $mp $cm -x $tag -x com.sun:auto-snapshot $debug -dF $target/$(echo $name | cut -d'/' -f1); then
                    # receive mounts the new filesystem despite canmount=noauto: unmount it, so it does not show
                    # up as mounted filesystem (e.g. in checkmk) until the next reboot
                    if [[ $fstype == "filesystem" ]] && [[ $(zfs get -H -o value mounted $target/$name) == "yes" ]]; then
                        debug "Unmounting $target/$name"
                        zfs unmount $target/$name
                    fi
                else
                    rc=1; log "[ERROR] initial replication to $target/$name failed."
                fi
            fi
            # the initial replication may take long, query the source snapshots again
            source_snaps=$(remote "zfs list -H -o name,guid -s creation -t snapshot $name 2>/dev/null")
        fi

        # incremental replication: from the newest target snapshot (matched by guid) to all newer source snapshots
        debug "$name - Start incremental replication"
        guid=$(zfs list -H -o guid -s creation -t snapshot $target/$name | tail -1)
        if [[ $guid != "" ]]; then
            last=$(grep $guid <<< "$source_snaps" | tail -1 | cut -f1)
            if [[ $last != "" ]]; then
                IFS=$'\n'
                if [[ $fstype == "filesystem" ]] && [[ $(zfs get -H -o value canmount $target/$name) != "noauto" ]]; then
                    zfs set canmount=noauto $target/$name
                fi
                for snap in $(grep -E "@.*($snapshot_filter)" <<< "$source_snaps" | grep --after-context=200 $guid | grep -v $guid | cut -f1); do
                    log "[INFO] Replicating delta of $last => $snap to $target/$name"
                    remote "zfs send -w $debug -i $last $snap" | zfs receive -x $tag -x com.sun:auto-snapshot -F $debug $target/$name \
                        || { rc=1; log "[ERROR] incremental replication to $target/$name failed."; }
                    last=$snap
                done
            else
                log "[ERROR] No matching snapshot (guid: $guid) found on source filesystem. This can be outdated snapshots on target or a previously deleted and new created dataset $name on the source filesystem."
                if [ $rc -eq 0 ]; then rc=1; fi
            fi
        else
            log "[ERROR] No snapshot found on $target/$name to add incremental snapshots to. The target dataset (with all children) needs to be deleted and recreated via replication."
            if [ $rc -eq 0 ]; then rc=1; fi
        fi

        # delete target snapshots older than the oldest source snapshot of each interval, keep at least min_keep
        debug "$name - Start deletion of old snapshots"
        # source snapshots (newest first), queried again after the replication
        source_snaps=$(remote "zfs list -H -o guid,name -S creation -t snapshot $name")
        filter=$(echo -e $snapshot_filter | sed "s/|/\n/g")
        IFS=$'\n'
        for interval in $filter ; do
            debug "$name - Checking interval $interval"
            guid=$(grep -E "@.*$interval" <<< "$source_snaps" | cut -f1 | tail -1)
            if [[ -n $guid ]]; then
                target_snaps=$(zfs list -H -o name,guid -S creation -t snapshot $target/$name | grep -E "@.*$interval")
                snaps_to_delete=$(grep --after-context=200 $guid <<< "$target_snaps" | grep -v $guid | cut -f1)
                snap_count=$(grep -c . <<< "$target_snaps")

                for snap in $snaps_to_delete; do
                    if [[ $snap_count -gt $min_keep ]]; then
                        log "[INFO] Deleting $snap"
                        debug "$name - snap_count=$snap_count, min_keep=$min_keep"
                        zfs destroy $debug $snap
                        ((snap_count--))
                    else
                        debug "$name - Skipping deletion of $snap. snap_count=$snap_count, min_keep=$min_keep"
                    fi
                done
            fi
        done
    else
        debug "$name - No snapshots found with filter $snapshot_filter"
    fi
    log "[INFO] Replication of $name finished."
done

#### monitoring: checkzfs output into the checkmk spool (local or on the source)
# the output goes to a temporary file of this run, is copied next to the spool file as dot file
# (ignored by the checkmk agent) and then renamed, so the agent never reads a partial file
if [ $checkzfs_disabled -eq 0 ]; then
    log "[INFO] Running checkzfs..."
    spooldir=/var/lib/check_mk_agent/spool
    spoolname=${checkzfs_spool_maxage}_${checkzfs_prefix}
    spoolfile=$(mktemp /tmp/.$spoolname.XXXXXX)
    spooltmp=$spooldir/$(basename $spoolfile)
    echo "<<<local>>>" > $spoolfile

    if [[ $checkzfs_local -eq 0 ]]; then
        c_key=--source
        c_value=$source:$sshport
    fi

    debug "Command:\n$checkzfs $c_key $c_value --output checkmk --threshold $checkzfs_max_age --maxsnapshots $checkzfs_max_snapshot_count --prefix $checkzfs_prefix --replicafilter "^${target}" --filter \"$checkzfs_filter\" >> $spoolfile"
    $checkzfs $c_key $c_value --output checkmk --threshold $checkzfs_max_age --maxsnapshots $checkzfs_max_snapshot_count --prefix $checkzfs_prefix --replicafilter "^${target}" --filter "$checkzfs_filter" >> $spoolfile
    # mktemp creates the file with mode 600, keep the previous mode of the spool file
    chmod 644 $spoolfile
    if [ $checkzfs_spool -eq 0 ] || [ $checkzfs_local -gt 0 ]; then
        mkdir -p $spooldir
        mv $spoolfile $spooltmp && mv $spooltmp $spooldir/$spoolname
    else
        remote "mkdir -p $spooldir"
        scp -P ${sshport} $spoolfile ${source}:$spooltmp && remote "mv $spooltmp $spooldir/$spoolname"
        rm -f $spoolfile
    fi
fi

exit $rc