#!/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 dir
    local -a logdirs=()
    # The caller passes its own application name so we archive the log of the
    # run that just finished, not whichever app happens to have left a log
    # lying around. Whitelist it: this value reaches us as a root-run argument
    # and is interpolated into a path, so reject anything unexpected.
    app="$1"
    case "$app" in
        iso-snapshot-cli|mx-snapshot) ;;
        *) return 1;;
    esac
    # Where the app wrote its log (see Log::defaultLogPath): a GUI run logs to
    # the caller's private runtime dir (/run/user/<uid>); a root CLI run logs to
    # /run. Both are non-world-writable, unlike the legacy /tmp fallback.
    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
    [[ "$uid" =~ ^[0-9]+$ && "$uid" -ne 0 ]] && logdirs+=("/run/user/${uid}")
    logdirs+=("/run" "/tmp")
    for dir in "${logdirs[@]}"; do
        [[ -d "$dir" ]] || continue
        src="${dir}/${app}.log"
        # Only copy an actual regular file, never a symlink, so the source
        # cannot be redirected at a root-owned file.
        [[ -f "$src" && ! -L "$src" ]] || continue
        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"
        return 0
    done
}

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";;
    datetime_log)
        datetime_log;;
    kill_mksquashfs)
        kill_mksquashfs;;
    drop_caches)
        drop_caches;;
esac
}

main "$@"
