#!/bin/sh
# Copyright 2022 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
verbose=false
syslog=false

# helper functions
GREEN=$'\033[0;32m'
RED=$'\033[0;31m'
CYAN=$'\033[;36m'
RESET=$'\033[0m'

print_help() {
            echo "${CYAN}[?] Usage:"
            echo "${0##*/} [-h, --help] [-t, --timeshift-auto] [-l, --log-file LOG_FILE] [-v, --verbose] [-s, --syslog] SNAPSHOTS_DIR"
            echo
            echo "SNAPSHOTS_DIR         Snapshot directory to watch, without effect when --timeshift-auto"
            echo
            echo "Optional arguments:"
            echo "-t, --timeshift-auto  Automatically detect Timeshifts snapshot directory"
            echo "-l, --log-file        Specify a logfile to write to"
            echo "-v, --verbose         Let the log of the daemon be more verbose"
            echo "-s, --syslog          Write to syslog"
            echo "-h, --help            Display this message${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:tvsh-: opt; do
    case "$opt" in
        -)
            case "${OPTARG}" in
                log-file)
                    logfile="${!OPTIND}"; OPTIND=$(( $OPTIND + 1 ))
                    ;;
                timeshift-auto)
                    timeshift_auto=true
                    ;;
                verbose)
                    verbose=true
                    ;;
                syslog)
                    syslog=true
                    ;;
                *)
                    if [ "$OPTERR" = 1 ] && [ "${optspec:0:1}" != ":" ]; then
                        err "[!] Unknown option --${OPTARG}" "${RED}" >&2
                        echo
                    fi
                    print_help
                    exit 1
                    ;;
            esac;;
        l)
            logfile="${OPTARG}"
            ;;
        t)
            timeshift_auto=true
            ;;
        v)
            verbose=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}"

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

vlog "Arguments:"
vlog "Snapshot directory: $snapshots"
vlog "Timestift autodetection: $timeshift_auto"
vlog "Logfile: $logfile"

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 [ -s "${GRUB_BTRFS_GRUB_DIRNAME:-/boot/grub}/grub-btrfs.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
}

# start the actual daemon
vlog "Snapshot dir watchtimeout: $watchtime"
vlog "Entering infinite while" "${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
            snapshots="/run/timeshift/${timeshift_pid}/backup/timeshift-btrfs/snapshots"
            log "detected running Timeshift at daemon startup, PID is: $timeshift_pid"
            vlog "new snapshots directory is $snapshots"
        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}')
                snapshots="/run/timeshift/${timeshift_pid}/backup/timeshift-btrfs/snapshots"
                log "detected Timeshift startup, PID is: $timeshift_pid" "${CYAN}"
                vlog "new snapshots directory is $snapshots" "${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 "$snapshots" ]; do
            # watch the actual snapshots folder for a new snapshot or a deletion of a snapshot
            if [ ${runs} = false ] && [ ${verbose} = false ]; then
                log "Watching $snapshots for new snapshots..." "${CYAN}"
            else
                vlog "Watching $snapshots for new snapshots..." "${CYAN}"
            fi
            runs=true
            inotifywait ${inotify_qiet_flag}  -e create -e delete -e unmount -t "$watchtime" "$snapshots" && {
                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 "${snapshots}" ] ; then # in case someone deletes the snapshots folder (in snapper mode) to prevent the while loop from going wild
        break
    fi
done

exit 0 # tradition is tradition