grub-btrfs/grub-btrfsd

#!/usr/bin/env bash
# Copyright 2023 Pascal Jaeger
# Distributed under the terms of the GNU General Public License v3
# Update GRUB when new BTRFS snapshots are created.

# init
timeshift_pid=-1
watchtime=0
logfile=0
snapshots=-1
timeshift_auto=false
timeshift_old=false
verbose=false
syslog=false
recursive=false

setcolors() {
    if [ "${1}" = true ]; then
          GREEN=$'\033[0;32m'
          RED=$'\033[0;31m'
          CYAN=$'\033[;36m'
          RESET=$'\033[0m'
    fi
    if [ "${1}" = false ]; then
        GREEN=$'\033[0;0m'
        RED=$'\033[0;0m'
        CYAN=$'\033[;0m'
        RESET=$'\033[0m'
    fi
}
setcolors true # normally we want colors

sysconfdir="/etc"
grub_btrfs_config="${sysconfdir}/default/grub-btrfs/config"
# source config file
[ -f "$grub_btrfs_config" ] && . "$grub_btrfs_config"
[ -f "${sysconfdir}/default/grub" ] && . "${sysconfdir}/default/grub"

print_help() {
            echo "${CYAN}[?] Usage:"
            echo "${0##*/} [-h, --help] [-c, --no-color] [-l, --log-file LOG_FILE] [-r, --recursive] [-s, --syslog] [-t, --timeshift-auto] [-v, --verbose] SNAPSHOTS_DIRS"
            echo
            echo "SNAPSHOTS_DIRS         Snapshot directories to watch, without effect when --timeshift-auto"
            echo
            echo "Optional arguments:"
            echo "-c, --no-color        Disable colors in output"
            echo "-l, --log-file        Specify a logfile to write to"
            echo "-r, --recursive       Watch snapshots directory recursively"
            echo "-s, --syslog          Write to syslog"
            echo "-o, --timeshift-old   Look for snapshots in directory of Timeshift <v22.06 (requires --timeshift-auto)"
            echo "-t, --timeshift-auto  Automatically detect Timeshifts snapshot directory"
            echo "-v, --verbose         Let the log of the daemon be more verbose"
            echo "-h, --help            Display this message"
            echo
            echo "Version ${GRUB_BTRFS_VERSION}${RESET}"
}

log() {
    echo "${2}"$1"${RESET}"
    if [ ${syslog} = true ]; then
        logger -p user.notice -t ${0##*/}"["$$"]" "$1"
    fi
    if [ ${#logfile} -gt 1  ]; then
        echo "$(date) ${1}" >> "${logfile}"
    fi
}

vlog() {
    if [ ${verbose} = true ]; then
       echo "${2}"$1"${RESET}"
       if [ ${syslog} = true ]; then
           logger -p user.notice -t ${0##*/} "$1"
       fi
       if [ ${#logfile} -gt 1  ]; then
        echo "$(date) ${1}" >> "${logfile}"
       fi
    fi
}

err() {
  echo "${2}"${1}"${RESET}" >&2
  if [ ${syslog} = true ]; then
      logger -p user.error -t ${0##*/} "$1"
  fi
  if [ ${#logfile} -gt 1  ]; then
      echo "$(date) error: ${1}" >> "${logfile}"
  fi
}

# parse arguments
while getopts :l:ctvrsh-: opt; do
    case "$opt" in
        -)
            case "${OPTARG}" in
                no-color)
                    setcolors false
                    ;;
                log-file)
                    logfile="${!OPTIND}"; OPTIND=$(( $OPTIND + 1 ))
                    ;;
                timeshift-auto)
                    timeshift_auto=true
                    ;;
                timeshift-old)
                    timeshift_old=true
                    ;;
                verbose)
                    verbose=true
                    ;;
                recursive)
                    recursive=true
                    ;;
                syslog)
                    syslog=true
                    ;;
                help)
                    print_help
                    exit 0
                    ;;
                *)
                    if [ "$OPTERR" = 1 ] && [ "${optspec:0:1}" != ":" ]; then
                        err "[!] Unknown option --${OPTARG}" "${RED}" >&2
                        echo
                    fi
                    print_help
                    exit 1
                    ;;
            esac;;
        c)
            setcolors false
            ;;
        l)
            logfile="${OPTARG}"
            ;;
        t)
            timeshift_auto=true
            ;;
        o)
            timeshift_old=true
            ;;
        v)
            verbose=true
            ;;
        r)
            recursive=true
            ;;
        s)
            syslog=true
            ;;
        h)
            print_help
            exit 0
            ;;
        *)
            if [ "$OPTERR" = 1 ] || [ "${optspec:0:1}" = ":" ]; then
                err "[!] Non-option argument: '-${OPTARG}'" "${RED}" >&2
                echo
            fi
            print_help
            exit 1
            ;;
    esac
done
shift $(( OPTIND - 1 ))

snapshots="${1}"

# check if inotify exists, see issue #227
if ! command -v inotifywait >/dev/null 2>&1; then
    err "[!] inotifywait was not found, exiting. Is inotify-tools installed?" "${RED}" >&2
    exit 1
fi

if [ ${#logfile} -gt 1  ]; then
    touch "${logfile}"
    echo "GRUB-BTRFSD log $(date)" >> "${logfile}"
fi

log  "grub-btrfsd starting up..." "${GREEN}"

if [ ${verbose} = true ]; then
    inotify_qiet_flag=""
else
    inotify_qiet_flag=" -q -q "
fi

if [ ${recursive} = true ]; then
    inotify_recursive_flag=" -r "
else
    inotify_recursive_flag=""
fi

if [ ${timeshift_auto} = false ] && [ ${timeshift_old} = true ]; then
    err "[!] Flag --timeshift-old requires flag --timeshift-auto" "${RED}" >&2
    exit 1
fi

vlog "Arguments:"
vlog "Snapshot directories: $snapshots"
vlog "Timestift autodetection: $timeshift_auto"
vlog "Timeshift old: $timeshift_old"
vlog "Logfile: $logfile"
vlog "Recursive: $recursive"

if ! [ -d "$snapshots" ] && ! [ ${timeshift_auto} = true ]; then
    err "[!] No directory found at ${snapshots}" "${RED}" >&2
    err "[!] Please specify a valid snapshot directory" "${RED}" >&2
    exit 1
fi

if [ ${timeshift_auto} = true ]; then
    watchtime=15
    [ -d /run/timeshift ] || mkdir /run/timeshift
else
    watchtime=0
fi

create_grub_menu() {
    #  create the grub submenu of the whole grub menu, depending on wether the submenu already exists
    #  and gives feedback if it worked
    if grep "snapshots-btrfs" "{grub_directory}/grub.cfg"; then
    	if  /etc/grub.d/41_snapshots-btrfs; then 
		log "Grub submenu recreated" "${GREEN}"
        else
		err "[!] Error during grub submenu creation (grub-btrfs error)" "${RED}"
        fi
    else
        if ${GRUB_BTRFS_MKCONFIG:-grub-mkconfig} -o "${GRUB_BTRFS_GRUB_DIRNAME:-/boot/grub}"/grub.cfg; then
		log "Grub menu recreated" "${GREEN}"
	else
		err "[!] Error during grub menu creation (grub/ grub-btrfs error)" "${RED}"
	fi
    fi
}

set_snapshot_dir() {
    # old timeshift has it's snapshot dir in a different location
    if [ "${timeshift_old}" = true ]; then
          snapdir="/run/timeshift/backup/timeshift-btrfs/snapshots"
    else
          snapdir="/run/timeshift/${timeshift_pid}/backup/timeshift-btrfs/snapshots"
    fi
}

daemon_function() {
    # start the actual daemon
    snapdir=$1
    vlog "Snapshot dir watchtimeout: $watchtime"
    vlog "Entering infinite while for $snapdir" "${GREEN}"
    while true; do
          runs=false
          if [ ${timeshift_auto} = true ] && ! [ "${timeshift_pid}" -gt 0 ] ; then
                # watch the timeshift folder for a folder that is created when timeshift starts up
                sleep 1 # for safety so the outer while is not going crazy
                if [ "${timeshift_pid}" -eq -2 ]; then
                      log "detected timeshift shutdown"
                fi
                timeshift_pid=$(ps ax | awk '{sub(/.*\//, "", $5)} $5 ~ /timeshift/ {print $1}')
                if [ "${#timeshift_pid}" -gt 0 ]; then
                      set_snapshot_dir
                      log "detected running Timeshift at daemon startup, PID is: $timeshift_pid"
                      vlog "new snapshots directory is $snapdir"
                else
                    log "Watching /run/timeshift for timeshift to start"
                    inotifywait ${inotify_qiet_flag} -e create -e delete /run/timeshift && {
                        sleep 1
                        timeshift_pid=$(ps ax | awk '{sub(/.*\//, "", $5)} $5 ~ /timeshift/ {print $1}')
                        set_snapshot_dir
                        log "detected Timeshift startup, PID is: $timeshift_pid" "${CYAN}"
                        vlog "new snapshots directory is $snapdir" "${CYAN}"
                        (create_grub_menu) # create the grub menu once immidiatly in a forking process. Snapshots from commandline using timeshift --create need this
                    }
                fi
                runs=false
          else
              while [ -d "$snapdir" ]; do
                    # watch the actual snapshots folder for a new snapshot or a deletion of a snapshot
                    if [ ${runs} = false ] && [ ${verbose} = false ]; then
                          log "Watching $snapdir for new snapshots..." "${CYAN}"
                    else
                        vlog "Watching $snapdir for new snapshots..." "${CYAN}"
                    fi
                    runs=true
                    inotifywait ${inotify_qiet_flag} ${inotify_recursive_flag} -e create -e delete -e unmount -t "$watchtime" "$snapdir" && {
                        log "Detected snapshot creation/ deletion, recreating Grub menu" "${CYAN}"
                        sleep 5
                        create_grub_menu
                    }
                    sleep 1
              done
              timeshift_pid=-2
          fi
          if ! [ ${timeshift_auto} = true ] && ! [ -d "${snapdir}" ] ; then # in case someone deletes the snapshots folder (in snapper mode) to prevent the while loop from going wild
                break
          fi
    done
}

# for all dirs that got passed to the script, start a new fork with that dir
snapdirs=( "${@}" )
for snapdir in "${snapdirs[@]}"
  do
      vlog "starting daemon with watching $snapdir..."
      daemon_function "${snapdir}" &
done

wait # wait for forks to finish, kill child forks if SIGTERM is sent


exit 0 # tradition is tradition