#!/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.

sighandler() {
    trap '""' SIGINT SIGTERM
    vlog "Parent $$: Received signal SIGINT/ SIGTERM"
    kill 0
    wait
    exit 0
}

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
}

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
}

checks() {
    # 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 [ ${timeshift_auto} = false ] && [ ${timeshift_old} = true ]; then
          err "[!] Flag --timeshift-old requires flag --timeshift-auto" "${RED}" >&2
          exit 1
    fi

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

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

    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} = 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 whether the submenu already exists
    #  and gives feedback if it worked
    if grep "snapshots-btrfs" "${GRUB_BTRFS_GRUB_DIRNAME:-/boot/grub}"/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() {
    trap 'vlog "$BASHPID: Received SIGINT/ SIGTERM"; exit 0' SIGINT SIGTERM
    # start the actual daemon
    snapdir=$1
    vlog "Subdaemon function started, PID: $BASHPID" "${GREEN}"
    vlog "${BASHPID}: Entering infinite while for $snapdir" "${GREEN}"
    vlog "${BASHPID}: Snapshot dir watchtimeout: $watchtime"
    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 "${BASHPID}: 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 "${BASHPID}: detected running Timeshift at daemon startup, PID is: $timeshift_pid"
                        vlog "${BASHPID}: 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 "${BASHPID}: detected Timeshift startup, PID is: $timeshift_pid" "${CYAN}"
                          vlog "${BASHPID}: new snapshots directory is $snapdir" "${CYAN}"
                          (create_grub_menu) # create the grub menu once immediately 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 "${BASHPID}: Watching $snapdir for new snapshots..." "${CYAN}"
                        else
                            vlog "${BASHPID}: 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 "${BASHPID}: 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
}

main() {
    # init
    timeshift_pid=-1
    watchtime=0
    logfile=0
    snapshots=-1
    timeshift_auto=false
    timeshift_old=false
    verbose=false
    syslog=false
    recursive=false
    trap sighandler SIGINT SIGTERM

    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"

    parse_arguments "${@}"
    shift $(( OPTIND - 1 ))
    snapdirs=( "${@}" )

    vlog "Arguments:"
    vlog "Snapshot directories: ${snapdirs[*]}"
    vlog "Timestift autodetection: $timeshift_auto"
    vlog "Timeshift old: $timeshift_old"
    vlog "Logfile: $logfile"
    vlog "Recursive: $recursive"

    checks
    setup

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

    if [ ${timeshift_auto} = true ] ; then
          daemon_function "timeshift" &
    else
        # for all dirs that got passed to the script, start a new fork with that dir
        for snapdir in "${snapdirs[@]}"
            do
                vlog "starting daemon watching $snapdir..."
                daemon_function "${snapdir}" &
        done
    fi
    wait # wait for forks to finish, kill child forks if SIGTERM is sent
    exit 0
}

main "${@}"