#!/bin/bash

# Root level functions requiring password for mx-snapshot

cleanup() {
    kill_mksquashfs
    # installed-to-live is Debian/MX-only (mx-remaster); on Arch the bind-root
    # is torn down by installed-to-live-arch instead, so a missing script is
    # not an error here — succeed as a no-op rather than exiting 127.
    if [[ -x /usr/sbin/installed-to-live ]]; then
        /usr/sbin/installed-to-live cleanup
    elif [[ -x /usr/bin/installed-to-live ]]; then
        /usr/bin/installed-to-live cleanup
    fi
}

# Unmount a tree with retries, falling back to a lazy detach. The lazy
# fallback matters for safety, not just tidiness: once detached, the mount is
# gone from the namespace, so a later rm -rf cannot traverse into it.
umount_retry() {
    local target=$1 i
    for ((i = 0; i < 10; i++)); do
        mountpoint -q "$target" || return 0
        umount --recursive "$target" 2>/dev/null
        sleep 0.3
    done
    umount --recursive --lazy "$target" 2>/dev/null
    ! mountpoint -q "$target"
}

cleanup_overlay() {
    local app_name overlay_base lower_dir
    app_name=$1
    [[ -z "$app_name" ]] && exit 2
    [[ "$app_name" =~ ^[A-Za-z0-9._-]+$ ]] || exit 2
    # Reject path-traversal: ".." would let overlay_base escape /run.
    [[ "$app_name" == *..* || "$app_name" == .* ]] && exit 2
    overlay_base="/run/${app_name}/bind-root-overlay"
    lower_dir="${overlay_base}/lower"
    local overlay_root="${overlay_base}/root"

    # Order matters: the overlay pins its lowerdir, so the merged root must go
    # first or the lower bind can never unmount.
    umount_retry "$overlay_root"
    umount_retry "$lower_dir"
    [[ -L "$overlay_base" ]] && exit 3
    # NEVER delete while anything is still mounted below: lower is a bind of
    # the real / — an rm -rf through a live bind would wipe the system.
    if findmnt --raw --noheadings --output TARGET \
            | grep -q -e "^${overlay_base}$" -e "^${overlay_base}/"; then
        echo "cleanup_overlay: mounts still active under ${overlay_base}; refusing to delete" >&2
        exit 4
    fi
    # --one-file-system is a second line of defense: overlay_base lives on the
    # /run tmpfs, so rm refuses to cross into any foreign mount it might meet.
    if [[ -d "$overlay_base" && "$overlay_base" == /run/*/bind-root-overlay ]]; then
        rm -rf --one-file-system "$overlay_base"
    fi
}

copy_log() {
    local app src dest uid allowed
    # The caller passes its own application name and the log file it wrote
    # (Log::getLog()), so we archive the run that just finished rather than
    # guessing: a root run under sudo logs to /run, while an older unprivileged
    # run in the same session may have left a log in /run/user/<uid>.
    # Whitelist the name: it reaches us as a root-run argument and is
    # interpolated into paths.
    app="$1"
    src="$2"
    case "$app" in
        iso-snapshot-cli|mx-snapshot) ;;
        *) return 1;;
    esac
    # Accept only locations Log::defaultLogPath() can produce for this app: the
    # invoking user's runtime dir (unprivileged GUI), /run (root run), or the
    # legacy /tmp fallback. Any other path is refused, so this root copy can
    # never be pointed at an arbitrary file.
    if [[ "$PKEXEC_UID" =~ ^[0-9]+$ ]]; then
        uid="$PKEXEC_UID"
    elif [[ "$SUDO_UID" =~ ^[0-9]+$ ]]; then
        uid="$SUDO_UID"
    else
        uid=$(id -u "$(logname)" 2>/dev/null)
    fi
    allowed=0
    [[ "$src" == "/run/${app}.log" || "$src" == "/tmp/${app}.log" ]] && allowed=1
    [[ "$uid" =~ ^[0-9]+$ && "$uid" -ne 0 && "$src" == "/run/user/${uid}/${app}.log" ]] && allowed=1
    [[ "$allowed" -eq 1 ]] || return 1
    # Only copy an actual regular file, never a symlink, so the source
    # cannot be redirected at a root-owned file.
    [[ -f "$src" && ! -L "$src" ]] || return 1
    dest="/var/log/${app}.log"
    [[ ! -e "$dest" && ! -L "$dest" ]] || {
        [[ -f "$dest" && ! -L "$dest" ]] || return 1
        mv -- "$dest" "${dest}.old"
    }
    # -P (no-dereference): if the file is swapped for a symlink after
    # the check (TOCTOU), copy the link itself, not its target.
    cp -P -- "$src" "$dest"
}

datetime_log() {
    date +"%Y%m%d_%H%M" > /etc/snapshot_created
}

# Kill mksquashfs and wait for it to actually exit. Cleanup unmounts the
# overlay right after this; a still-dying mksquashfs keeps the merged root
# busy, the busy overlay pins the lower bind of /, and teardown then has to
# leave mounts behind. Escalate to SIGKILL if TERM isn't enough.
kill_mksquashfs() {
    local i
    pgrep -x mksquashfs >/dev/null || return 0
    pkill -x mksquashfs
    for ((i = 0; i < 50; i++)); do
        pgrep -x mksquashfs >/dev/null || return 0
        sleep 0.1
    done
    pkill -9 -x mksquashfs
    for ((i = 0; i < 20; i++)); do
        pgrep -x mksquashfs >/dev/null || return 0
        sleep 0.1
    done
    return 1
}

drop_caches() {
    echo 1 > /proc/sys/vm/drop_caches
}

chown_conf() {
    local uid user home app file real

    # Determine the calling (non-root) user reliably. pkexec exports
    # PKEXEC_UID; fall back to logname for non-pkexec invocations.
    if [[ -n "$PKEXEC_UID" ]]; then
        uid="$PKEXEC_UID"
    else
        uid=$(id -u "$(logname)" 2>/dev/null)
    fi
    [[ "$uid" =~ ^[0-9]+$ ]] || exit 2
    # Never operate on behalf of root.
    [[ "$uid" -eq 0 ]] && exit 2

    user=$(getent passwd "$uid" | cut -d: -f1)
    home=$(getent passwd "$uid" | cut -d: -f6)
    [[ -z "$user" || -z "$home" || ! -d "$home" ]] && exit 2
    # Resolve the home path too, so the confinement check below still
    # matches when /home (or a parent) is itself a symlink.
    home=$(realpath -e "$home" 2>/dev/null) || exit 2

    for app in iso-snapshot-cli mx-snapshot; do
        file="${home}/.config/MX-Linux/${app}.conf"
        # Must be an existing regular file that is NOT itself a symlink.
        [[ -f "$file" && ! -L "$file" ]] || continue
        # Resolve symlinks in any path component and confine the result
        # to the user's home, defeating symlink-redirection attacks.
        real=$(realpath -e "$file" 2>/dev/null) || continue
        [[ "$real" == "$home"/* ]] || continue
        # -h: do not dereference, guarding against a TOCTOU symlink swap.
        chown -h "${user}:" "$real"
    done
}

main() {
case "$1" in
    chown_conf)
        chown_conf;;
    cleanup)
        cleanup;;
    cleanup_overlay)
        cleanup_overlay "$2";;
    copy_log)
        copy_log "$2" "$3";;
    datetime_log)
        datetime_log;;
    kill_mksquashfs)
        kill_mksquashfs;;
    drop_caches)
        drop_caches;;
esac
}

main "$@"
