#!/bin/bash
# Privileged helper for Lite Distro Builder.
# Invoked via pkexec; the polkit action com.linuxliteos.litedistrobuilder
# pins this path so polkit matches a single policy for every sub-action.
#
# Usage:
#   LiteDistroBuilder-helper dist <tmp_conf>
#   LiteDistroBuilder-helper clean <tmp_conf>
#   LiteDistroBuilder-helper skelcopy
#   LiteDistroBuilder-helper mkworkdir <path>
#   LiteDistroBuilder-helper plymouth-set <theme.plymouth>
#   LiteDistroBuilder-helper plymouth-auto
#   LiteDistroBuilder-helper plymouth-preview
#   LiteDistroBuilder-helper plymouth-install <theme_name> <theme_name_fixed> <source_image>
#   LiteDistroBuilder-helper plymouth-delete <theme_name_fixed>
#
# Security notes:
#   Arguments are validated to reject traversal and shell metacharacters; we
#   only accept bare alphanumerics + a few safe punctuation chars for theme
#   names and only absolute paths for file arguments.

set -e

PLYMOUTH_TEMPLATE="/etc/LiteDistroBuilder/plymouth/LiteDistroBuilder-theme"
THEMES_DIR="/usr/share/plymouth/themes"

_require_safe_name() {
    # Allow letters, digits, dash, underscore, dot. Reject /, .., whitespace.
    case "$1" in
        ''|*..*|*/*|*' '*|*$'\t'*) echo "Invalid name: $1" >&2; exit 2 ;;
    esac
    case "$1" in
        *[!A-Za-z0-9._-]*) echo "Invalid characters in name: $1" >&2; exit 2 ;;
    esac
}

_require_existing_file() {
    [ -n "$1" ] && [ -f "$1" ] || { echo "Not a file: $1" >&2; exit 2; }
}

_install_conf() {
    # Take a user-written temp conf and install to /etc/.
    local src="$1"
    [ -n "$src" ] || return 0
    _require_existing_file "$src"
    install -m 644 "$src" /etc/LiteDistroBuilder.conf
    rm -f "$src"
}

_ensure_workdir() {
    # Read WORKDIR out of the now-installed conf and create if missing.
    local workdir
    workdir="$(awk -F'=' '/^WORKDIR=/ {gsub(/"/,"",$2); print $2; exit}' /etc/LiteDistroBuilder.conf 2>/dev/null || true)"
    [ -n "$workdir" ] || workdir="/home/LiteDistroBuilder"
    case "$workdir" in /*) ;; *) echo "WORKDIR must be absolute: $workdir" >&2; exit 2 ;; esac
    mkdir -p "$workdir"
}

case "$1" in
    dist)
        _install_conf "$2"
        _ensure_workdir
        exec /usr/bin/LiteDistroBuilder dist
        ;;
    clean)
        _install_conf "$2"
        _ensure_workdir
        exec /usr/bin/LiteDistroBuilder clean
        ;;
    skelcopy)
        exec /usr/bin/LiteDistroBuilder-skelcopy-home
        ;;
    cancel)
        # Cancel a running build. Arg 2 is an optional PGID from the GUI's
        # tcgetpgrp() on the VTE — if supplied we SIGTERM that group first.
        # Either way, sweep any lingering LiteDistroBuilder processes as a
        # belt-and-braces in case the PGID lookup raced with process exit.
        pgid="$2"
        if [ -n "$pgid" ] && [ "$pgid" -gt 0 ] 2>/dev/null; then
            kill -TERM -"$pgid" 2>/dev/null || true
        fi
        pkill -TERM -x LiteDistroBuilder              2>/dev/null || true
        pkill -TERM -x LiteDistroBuilder-skelcopy-home 2>/dev/null || true
        pkill -TERM -x mksquashfs                      2>/dev/null || true
        pkill -TERM -x xorriso                         2>/dev/null || true
        exit 0
        ;;
    mkworkdir)
        [ -n "$2" ] || exit 2
        case "$2" in /*) ;; *) echo "Must be absolute: $2" >&2; exit 2 ;; esac
        mkdir -p "$2"
        [ -n "$PKEXEC_UID" ] && chown "$PKEXEC_UID" "$2" 2>/dev/null || true
        ;;
    plymouth-set)
        _require_existing_file "$2"
        /usr/bin/update-alternatives --set default.plymouth "$2"
        # Plymouth themes live in the initramfs — a bare `update-alternatives`
        # changes the config pointer but the *running* or *next-boot* initrd
        # still contains the previous theme. Rebuild so the selection actually
        # applies. ~20s on a typical box; the GUI pulses a progress bar.
        kver="$(uname -r)"
        exec /usr/sbin/mkinitramfs -o "/boot/initrd.img-$kver" "$kver"
        ;;
    plymouth-auto)
        /usr/bin/update-alternatives --auto default.plymouth
        # Same rebuild story as plymouth-set — the initramfs is the thing
        # that actually ships the theme to the boot sequence.
        kver="$(uname -r)"
        exec /usr/sbin/mkinitramfs -o "/boot/initrd.img-$kver" "$kver"
        ;;
    plymouth-preview)
        exec /usr/bin/plymouth-preview
        ;;
    plymouth-install)
        # $2 = human-readable theme_name (may contain spaces)
        # $3 = sanitized theme_name_fixed  (safe filesystem name)
        # $4 = source image (absolute path)
        local_theme_name="$2"
        local_theme_name_fixed="$3"
        local_source="$4"
        [ -n "$local_theme_name" ] || { echo "theme_name required" >&2; exit 2; }
        _require_safe_name "$local_theme_name_fixed"
        _require_existing_file "$local_source"

        theme_dir="$THEMES_DIR/$local_theme_name_fixed"
        rm -rf "$theme_dir"
        mkdir -p "$theme_dir"

        img_basename="$(basename "$local_source")"
        cp "$local_source" "$theme_dir/$img_basename"
        cp "$PLYMOUTH_TEMPLATE/progress_bar.png" "$theme_dir/"
        cp "$PLYMOUTH_TEMPLATE/progress_box.png" "$theme_dir/"

        # Script: substitute placeholder for the image filename.
        awk -v pic="$img_basename" '{gsub(/__THEMEPIC__/, pic); print}' \
            "$PLYMOUTH_TEMPLATE/LiteDistroBuilder-theme.script" \
            > "$theme_dir/$local_theme_name_fixed.script"

        # Plymouth config: substitute theme name + directory.
        awk -v name="$local_theme_name" -v dir="$local_theme_name_fixed" \
            '{gsub(/__THEMENAME__/, name); gsub(/__THEMEDIR__/, dir); print}' \
            "$PLYMOUTH_TEMPLATE/LiteDistroBuilder-theme.plymouth" \
            > "$theme_dir/$local_theme_name_fixed.plymouth"

        /usr/bin/update-alternatives --install \
            /usr/share/plymouth/themes/default.plymouth default.plymouth \
            "$theme_dir/$local_theme_name_fixed.plymouth" 80
        /usr/bin/update-alternatives --set default.plymouth \
            "$theme_dir/$local_theme_name_fixed.plymouth"

        kver="$(uname -r)"
        /usr/sbin/mkinitramfs -o "/boot/initrd.img-$kver" "$kver"
        ;;
    plymouth-delete)
        _require_safe_name "$2"
        target="$THEMES_DIR/$2"
        # Belt-and-braces: confirm the resolved path is still inside THEMES_DIR.
        case "$target" in "$THEMES_DIR"/*) ;; *) echo "Refusing to delete $target" >&2; exit 2 ;; esac
        rm -rf "$target"
        ;;
    *)
        echo "Unknown action: $1" >&2
        exit 2
        ;;
esac
