#!/bin/bash
# computle-kiosk-setup enable|disable: makes this device a Computle kiosk terminal (account, autologin into Cage, silent branded boot), or puts it back.
# computle-kioskd runs it when the portal pairs or removes the terminal; an administrator can run it by hand. Both directions are idempotent.
set -uo pipefail

KIOSK_USER=computle-kiosk
FULL_NAME=Computle
DATA=${COMPUTLE_KIOSK_DATA:-/usr/share/computle-stream/kiosk}
STATE=/var/lib/computle-kiosk-setup
UNIT=computle-kiosk.service
UNIT_FILE=/etc/systemd/system/$UNIT
GETTY_DROPIN=/etc/systemd/system/getty@tty1.service.d/50-computle-kiosk.conf
LOGIND_DROPIN=/etc/systemd/logind.conf.d/50-computle-kiosk-vt.conf
POLKIT_RULE=/etc/polkit-1/rules.d/50-computle-kiosk.rules
SYSCTL_FILE=/etc/sysctl.d/90-computle-kiosk.conf
GRUB_DROPIN=/etc/default/grub.d/99-computle-kiosk.cfg
ARMBIAN_ENV=/boot/armbianEnv.txt
THEME=computle-os
THEME_DIR=/usr/share/plymouth/themes/$THEME
QUIET_ARGS="quiet splash loglevel=3 vt.global_cursor_default=0 plymouth.ignore-serial-consoles systemd.show_status=false rd.udev.log_level=3 udev.log_priority=3"

changed=0
warnings=0
initramfs=0

log() { echo "computle-kiosk-setup: $*"; }
warn() {
    echo "computle-kiosk-setup: warning: $*" >&2
    warnings=$((warnings + 1))
}
die() {
    echo "computle-kiosk-setup: $*" >&2
    exit "${2:-1}"
}

# Writes stdin to path only when it differs; returns 0 when it changed the file.
put() {
    local path=$1 mode=$2 temp
    temp=$(mktemp) || die "cannot create a temporary file"
    cat >"$temp"
    if [ -f "$path" ] && cmp -s "$temp" "$path"; then
        rm -f "$temp"
        return 1
    fi
    mkdir -p "$(dirname "$path")"
    install -m "$mode" -o root -g root "$temp" "$path" || die "cannot write $path"
    rm -f "$temp"
    changed=1
    log "wrote $path"
    return 0
}

put_file() { put "$2" "$3" <"$1"; }

drop() {
    local path=$1
    [ -e "$path" ] || return 1
    rm -f "$path"
    changed=1
    log "removed $path"
    return 0
}

remember() {
    local name=$1 value=$2
    mkdir -p "$STATE"
    chmod 700 "$STATE"
    [ -e "$STATE/$name" ] || printf '%s\n' "$value" >"$STATE/$name"
}

recalled() { [ -e "$STATE/$1" ] && cat "$STATE/$1"; }

have() { command -v "$1" >/dev/null 2>&1; }

systemd_running() { [ -d /run/systemd/system ]; }

# ---- the kiosk account ----

ensure_user() {
    local gecos
    if gecos=$(getent passwd "$KIOSK_USER" | cut -d: -f5); then
        case "$gecos" in
        "$FULL_NAME"*) ;;
        *) die "a $KIOSK_USER account exists that this tool did not make ($gecos); left alone" 3 ;;
        esac
    else
        useradd --create-home --comment "$FULL_NAME" --shell /usr/sbin/nologin --user-group "$KIOSK_USER" || die "cannot create $KIOSK_USER"
        changed=1
        log "created the $KIOSK_USER account"
    fi
    passwd --lock "$KIOSK_USER" >/dev/null 2>&1 || true
    local group
    for group in video audio input render plugdev netdev bluetooth; do
        getent group "$group" >/dev/null || continue
        if ! id -nG "$KIOSK_USER" | tr ' ' '\n' | grep -qx "$group"; then
            usermod -a -G "$group" "$KIOSK_USER" && changed=1
        fi
    done
    local home
    home=$(getent passwd "$KIOSK_USER" | cut -d: -f6)
    if [ -n "$home" ] && [ -d "$home" ] && [ ! -e "$home/.hushlogin" ]; then
        touch "$home/.hushlogin"
        chown "$KIOSK_USER:" "$home/.hushlogin"
    fi
}

remove_user() {
    local gecos
    gecos=$(getent passwd "$KIOSK_USER" | cut -d: -f5) || return 0
    case "$gecos" in
    "$FULL_NAME"*) ;;
    *) warn "$KIOSK_USER was not made by this tool; left alone"; return 0 ;;
    esac
    have loginctl && loginctl terminate-user "$KIOSK_USER" >/dev/null 2>&1
    userdel --remove "$KIOSK_USER" >/dev/null 2>&1 || userdel "$KIOSK_USER" || warn "cannot remove $KIOSK_USER"
    changed=1
    log "removed the $KIOSK_USER account"
}

# ---- session: Cage on tty1 in place of the display manager ----

current_display_manager() {
    local target
    target=$(readlink /etc/systemd/system/display-manager.service 2>/dev/null) || return 0
    basename "$target"
}

enable_session() {
    put_file "$DATA/$UNIT" "$UNIT_FILE" 0644
    printf '[Unit]\n# tty1 belongs to the kiosk session; its login prompt appears only during maintenance.\nConditionPathExists=/run/computle/kiosk-maintenance\n' | put "$GETTY_DROPIN" 0644
    printf '[Login]\nNAutoVTs=0\nReserveVT=0\n' | put "$LOGIND_DROPIN" 0644
    put_file "$DATA/50-computle-kiosk.rules" "$POLKIT_RULE" 0644

    local dm
    dm=$(current_display_manager)
    if [ -n "$dm" ] && [ "$dm" != "$UNIT" ]; then
        remember display-manager "$dm"
        systemctl disable "$dm" >/dev/null 2>&1 || warn "cannot disable $dm"
        changed=1
        log "the display manager $dm no longer starts; the kiosk session does"
    fi
    if [ "$(current_display_manager)" != "$UNIT" ] || ! systemctl is-enabled --quiet "$UNIT" 2>/dev/null; then
        systemctl enable --force "$UNIT" >/dev/null 2>&1 || die "cannot enable $UNIT"
        changed=1
    fi
    local target
    target=$(systemctl get-default 2>/dev/null)
    if [ "$target" != graphical.target ]; then
        remember default-target "$target"
        systemctl set-default graphical.target >/dev/null 2>&1 || warn "cannot make graphical.target the default"
        changed=1
    fi
    systemctl enable computle-kioskd.service >/dev/null 2>&1 || warn "cannot enable computle-kioskd.service"
    if systemd_running; then
        systemctl daemon-reload || warn "systemctl daemon-reload failed"
    fi
}

disable_session() {
    if systemctl is-enabled --quiet "$UNIT" 2>/dev/null; then
        systemctl disable "$UNIT" >/dev/null 2>&1 || warn "cannot disable $UNIT"
        changed=1
    fi
    drop "$UNIT_FILE"
    drop "$GETTY_DROPIN"
    rmdir "$(dirname "$GETTY_DROPIN")" 2>/dev/null
    drop "$LOGIND_DROPIN"
    drop "$POLKIT_RULE"
    local dm target
    if dm=$(recalled display-manager); then
        systemctl enable --force "$dm" >/dev/null 2>&1 || warn "cannot enable $dm again"
        rm -f "$STATE/display-manager"
        changed=1
        log "the display manager $dm starts again"
    fi
    if target=$(recalled default-target); then
        [ -n "$target" ] && systemctl set-default "$target" >/dev/null 2>&1
        rm -f "$STATE/default-target"
    fi
    if systemd_running; then
        systemctl daemon-reload || warn "systemctl daemon-reload failed"
    fi
}

# ---- Plymouth: the Computle OS theme ----

current_theme() {
    if have plymouth-set-default-theme; then
        plymouth-set-default-theme 2>/dev/null
    elif [ -e /usr/share/plymouth/themes/default.plymouth ]; then
        basename "$(dirname "$(readlink -f /usr/share/plymouth/themes/default.plymouth)")"
    fi
}

enable_theme() {
    if [ ! -d /usr/share/plymouth/themes ]; then
        warn "Plymouth is not installed; the boot shows no Computle logo"
        return
    fi
    local file
    for file in computle-os.plymouth computle-os.script logo.png; do
        put_file "$DATA/plymouth/$THEME/$file" "$THEME_DIR/$file" 0644 && initramfs=1
    done
    local theme
    theme=$(current_theme)
    if [ "$theme" != "$THEME" ]; then
        remember plymouth-theme "$theme"
        if have update-alternatives && [ -e /usr/share/plymouth/themes/default.plymouth ]; then
            remember plymouth-alternative "$(readlink -f /usr/share/plymouth/themes/default.plymouth)"
            update-alternatives --quiet --install /usr/share/plymouth/themes/default.plymouth default.plymouth "$THEME_DIR/$THEME.plymouth" 200
            update-alternatives --quiet --set default.plymouth "$THEME_DIR/$THEME.plymouth"
        fi
        if have plymouth-set-default-theme; then
            plymouth-set-default-theme "$THEME" || warn "cannot select the $THEME theme"
        fi
        initramfs=1
        changed=1
        log "the boot splash is now Computle OS"
    fi
}

disable_theme() {
    [ -d "$THEME_DIR" ] || return
    local previous alternative
    previous=$(recalled plymouth-theme)
    alternative=$(recalled plymouth-alternative)
    if have update-alternatives; then
        update-alternatives --quiet --remove default.plymouth "$THEME_DIR/$THEME.plymouth" 2>/dev/null
        if [ -n "$alternative" ] && [ -e "$alternative" ]; then
            update-alternatives --quiet --set default.plymouth "$alternative" 2>/dev/null || warn "cannot select $alternative again"
        fi
    fi
    if [ -n "$previous" ] && have plymouth-set-default-theme; then
        plymouth-set-default-theme "$previous" 2>/dev/null || warn "cannot select the $previous theme again"
    elif have plymouth-set-default-theme && [ "$(plymouth-set-default-theme 2>/dev/null)" = "$THEME" ]; then
        plymouth-set-default-theme --reset 2>/dev/null
    fi
    rm -f "$STATE/plymouth-theme" "$STATE/plymouth-alternative"
    rm -rf "$THEME_DIR"
    initramfs=1
    changed=1
    log "the boot splash is back to ${previous:-the default}"
}

rebuild_initramfs() {
    [ "$initramfs" = 1 ] || return
    if have update-initramfs; then
        update-initramfs -u >/dev/null 2>&1 || warn "update-initramfs failed; the splash changes at the next kernel update"
    elif have dracut; then
        dracut --force >/dev/null 2>&1 || warn "dracut failed; the splash changes at the next kernel update"
    fi
}

# ---- silent boot: GRUB (x86) and U-Boot through armbianEnv.txt (Armbian) ----

enable_grub() {
    [ -f /etc/default/grub ] || return
    # recordfail would show the menu for 30 s after every power cut, which is how kiosks are often switched off.
    if cat <<EOF | put "$GRUB_DROPIN" 0644; then
GRUB_TIMEOUT=0
GRUB_TIMEOUT_STYLE=hidden
GRUB_RECORDFAIL_TIMEOUT=0
GRUB_DISABLE_OS_PROBER=true
GRUB_TERMINAL=console
GRUB_GFXPAYLOAD_LINUX=keep
GRUB_CMDLINE_LINUX_DEFAULT="\${GRUB_CMDLINE_LINUX_DEFAULT} $QUIET_ARGS"
EOF
        update_grub
    fi
}

disable_grub() {
    drop "$GRUB_DROPIN" && update_grub
}

update_grub() {
    if have update-grub; then
        update-grub >/dev/null 2>&1 || warn "update-grub failed; the boot menu changes the next time it runs"
    elif have grub2-mkconfig; then
        grub2-mkconfig -o /boot/grub2/grub.cfg >/dev/null 2>&1 || warn "grub2-mkconfig failed"
    else
        warn "no update-grub; the boot menu changes the next time GRUB's configuration is generated"
    fi
}

# One key=value line of armbianEnv.txt, or nothing.
armbian_get() { grep -m1 "^$1=" "$ARMBIAN_ENV" | cut -d= -f2-; }

armbian_set() {
    local key=$1 value=$2 temp
    temp=$(mktemp)
    if grep -q "^$key=" "$ARMBIAN_ENV"; then
        awk -v k="$key" -v v="$value" 'BEGIN { FS = OFS = "=" } $1 == k { print k "=" v; next } { print }' "$ARMBIAN_ENV" >"$temp"
    else
        cat "$ARMBIAN_ENV" >"$temp"
        printf '%s=%s\n' "$key" "$value" >>"$temp"
    fi
    if ! cmp -s "$temp" "$ARMBIAN_ENV"; then
        cat "$temp" >"$ARMBIAN_ENV"
        changed=1
    fi
    rm -f "$temp"
}

armbian_unset() {
    local key=$1 temp
    grep -q "^$key=" "$ARMBIAN_ENV" || return
    temp=$(mktemp)
    grep -v "^$key=" "$ARMBIAN_ENV" >"$temp"
    cat "$temp" >"$ARMBIAN_ENV"
    rm -f "$temp"
    changed=1
}

# U-Boot's verbosity and logo, and the kernel's console on the debug UART rather than the screen, which then shows only the splash.
enable_armbian() {
    [ -f "$ARMBIAN_ENV" ] || return
    local key before
    before=$(cat "$ARMBIAN_ENV")
    if [ ! -e "$STATE/armbianEnv.orig" ]; then
        mkdir -p "$STATE"
        chmod 700 "$STATE"
        for key in verbosity bootlogo console extraargs; do
            if grep -q "^$key=" "$ARMBIAN_ENV"; then
                grep -m1 "^$key=" "$ARMBIAN_ENV"
            else
                printf '%s\n' "-$key"
            fi
        done >"$STATE/armbianEnv.orig"
    fi
    armbian_set verbosity 0
    armbian_set bootlogo true
    armbian_set console serial
    local extra arg
    extra=$(armbian_get extraargs)
    for arg in $QUIET_ARGS; do
        case " $extra " in
        *" $arg "*) ;;
        *) extra="${extra:+$extra }$arg" ;;
        esac
    done
    armbian_set extraargs "$extra"
    [ "$before" != "$(cat "$ARMBIAN_ENV")" ] && log "armbianEnv.txt: quiet boot, Armbian's boot logo, kernel console on serial"
}

disable_armbian() {
    [ -f "$ARMBIAN_ENV" ] && [ -f "$STATE/armbianEnv.orig" ] || return
    local line
    while IFS= read -r line; do
        case "$line" in
        -*) armbian_unset "${line#-}" ;;
        *=*) armbian_set "${line%%=*}" "${line#*=}" ;;
        esac
    done <"$STATE/armbianEnv.orig"
    rm -f "$STATE/armbianEnv.orig"
    log "armbianEnv.txt is back as it was"
}

enable_kernel_quiet() {
    if printf 'kernel.printk = 3 3 3 3\n' | put "$SYSCTL_FILE" 0644; then
        sysctl -q -p "$SYSCTL_FILE" >/dev/null 2>&1 || warn "the console log level applies from the next start"
    fi
}

# ---- main ----

usage() {
    echo "usage: computle-kiosk-setup enable|disable [--remove-user] [--from-kioskd]" >&2
    exit 2
}

[ "$#" -ge 1 ] || usage
action=$1
shift
remove_account=0
from_kioskd=0
for arg in "$@"; do
    case "$arg" in
    --remove-user) remove_account=1 ;;
    --from-kioskd) from_kioskd=1 ;;
    *) usage ;;
    esac
done
[ "$(id -u)" = 0 ] || die "run as root" 4
[ -d "$DATA" ] || die "$DATA is missing; is computle-stream installed?" 5

case "$action" in
enable)
    ensure_user
    enable_session
    enable_theme
    rebuild_initramfs
    enable_grub
    enable_armbian
    enable_kernel_quiet
    if [ "$changed" = 1 ]; then
        [ "$from_kioskd" = 1 ] || log "done; restart the device to start kiosk mode"
    else
        log "already a kiosk terminal; nothing changed"
    fi
    ;;
disable)
    disable_session
    disable_theme
    rebuild_initramfs
    disable_grub
    disable_armbian
    drop "$SYSCTL_FILE"
    [ "$remove_account" = 1 ] && remove_user
    rmdir "$STATE" 2>/dev/null
    if [ "$changed" = 1 ]; then
        [ "$from_kioskd" = 1 ] || log "done; the device starts as it did before kiosk mode"
    else
        log "not a kiosk terminal; nothing changed"
    fi
    ;;
*)
    usage
    ;;
esac
[ "$warnings" = 0 ] || log "$warnings warning(s); see above"
exit 0
