#!/bin/bash

# SPDX-License-Identifier: GPL-2.0

# Copyright © Sébastien Luttringer
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

# Check running systemd services for binary update
# Convenient way to restart updated systemd service after upgrade

# bash options
shopt -s xpg_echo

# disable grep options to avoid non default behaviour
unset GREP_OPTIONS

# systemd cgroup path
if [[ -d /sys/fs/cgroup/systemd ]]; then
  SYSTEMD_CGROUP_BASE_PATH='/sys/fs/cgroup/systemd'
else
  SYSTEMD_CGROUP_BASE_PATH='/sys/fs/cgroup'
fi

# colors
if [[ -t 1 ]]; then
    C_BOLD='\e[1m'
    C_BLUE='\e[34m'
    C_RED='\e[31m'
    C_YELLOW='\e[33m'
    C_RESET='\e[m'
fi

# default options
AUTOCONFIRM=0       # autoconfirmation
FAILED=1            # display failed service at the end
PACDIFF=1           # run pacdiff
THREEWAY=0          # run pacdiff with --threeway
RELOAD=1            # reload systemd
RESTART=1           # restart services
SERIALIZE=0         # run in parallel
STATUS=1            # display status after systemctl
USER_SLICE=0        # act on users services
MACHINE_SLICE=0     # act on machine services

# ignored service list
IGNORED_SERVICES=(
    "getty@tty.*.service"
    "systemd-logind.service"
    "dbus-broker.service"
    "auditd.service"
)

# print $* as an arrow line
arrow() {
    printf "${C_BOLD}${C_BLUE}:: ${C_RESET}${C_BOLD}%s${C_RESET}\n" "$*"
}

# print $* as an warning message
warn() {
    printf "${C_BOLD}${C_YELLOW}Warning:: ${C_RESET}${C_BOLD}%s${C_RESET}\n" "$*" >&2
}

# print $* as an error message
error() {
    printf "${C_BOLD}${C_RED}Error:: ${C_RESET}${C_BOLD}%s${C_RESET}\n" "$*" >&2
}

# usage : in_array( $needle, $haystack )
# return : 0 - found
#          1 - not found
in_array() {
    local needle="$1"; shift
    local item
    for item in "$@"; do
       [[ $item = "$needle" ]] && return 0 # Found
    done
    return 1 # Not Found
}

# ask for confirmation
# return 0 when confirmed, otherwise 1
confirm() {
    (( AUTOCONFIRM == 1 )) && return 0
    local -i try
    local ans
    for ((try=1; try<=5; try++)); do
        printf '%s [Yes|No] ' "$1"
        read -r ans || return 1
        case $ans in
            y|Y|yes|Yes) return 0;;
            n|N|no|No) return 1;;
        esac
    done
    error "Too much invalid answer. Not confirmed."
    return 1
}

# get running systemd services
get_services() {
    local -a grep_patterns
    read -r -a grep_patterns <<< "$(printf -- '-e %s ' "${IGNORED_SERVICES[@]}")"
    systemctl --no-legend --full --type service --state running | tr -d '●' | awk '{print $1}' | grep -v "${grep_patterns[@]}"
}

# get systemd services with updated mapped files
get_broken_maps() {
    local service path pidfile unit_path maps_path pids deleted
    local -a pids=()
    local -i pid=0
    for service in $(get_services); do
        unit_path="$(systemctl --property ControlGroup --value show "$service")"
        # hack to fix to systemd internal cgroup escaping on slice
        unit_path="$(printf '%s' "$unit_path"|sed 's,\\x5c,\\,')"
        # has Delegate Subgroup?
        delegate_path="$(systemctl --property DelegateSubgroup --value show "$service")"
        [[ -n ${delegate_path} ]] && unit_path="${unit_path}/${delegate_path}"
        # get the right pidfile name
        pidfile=''
        for path in "$SYSTEMD_CGROUP_BASE_PATH$unit_path/cgroup.procs" \
            "$SYSTEMD_CGROUP_BASE_PATH$unit_path/tasks"; do
            [[ -r "$path" ]] && pidfile="$path" && continue
        done
        [[ -z "$pidfile" ]] && error "Unable to find pid file for $service." && continue
        # skip non system units
        (( USER_SLICE == 0 )) && [[ "$unit_path" =~ /user\.slice/ ]] && continue
        (( MACHINE_SLICE == 0 )) && [[ "$unit_path" =~ /machine\.slice/ ]] && continue
        # parse pidfile
        mapfile -t pids < "$pidfile"
        if (( "${#pids[*]}" == 0 )); then
            error "Unable to parse pid file for $service."
            continue
        fi
        for pid in "${pids[@]}"; do
            maps_path="/proc/$pid/maps"
            [[ -r "$maps_path" ]] || {
                error "Unable to read maps file of $service for pid $pid."
                continue
            }
            # only file mapped as executable
            deleted="$(grep -F '(deleted)' "$maps_path"|sed -nr 's|^\S+ ..x. \S+ \S+ \S+ \s+||p'|grep -v "/memfd:")"
            if [[ $deleted ]]; then
                printf "%s\n" "$service"
                break
            fi
        done
    done
}

# get dbus clients on the system bus
get_dbus_names() {
    dbus-send --system --dest=org.freedesktop.DBus --type=method_call --print-reply \
    /org/freedesktop/DBus org.freedesktop.DBus.ListNames|sed -rn 's/\s*string "(.*)"/\1/p'
}

# get systemd services not registered on dbus system bus
get_missing_dbus() {
    local service busname
    local -a registered
    mapfile -t registered < <(get_dbus_names)
    for service in $(get_services); do
        # get the service registered bus name
        busname="$(systemctl --property BusName --value show "$service")"
        if [[ "$busname" ]] && ! in_array "$busname" "${registered[@]}"; then
            echo "$service"
        fi
    done
}

# display restart instruction from service name
display_restart() {
    local service
    echo '-------8<-------------------------------8<---------'
    for service; do
        echo "systemctl restart '$service'"
    done
    echo '-------8<-------------------------------8<---------'
}

# restart systemd services given in arguments
restart_services() {
    local service
    local -i last_registered_pids_count
    local -A registered_pids=()
    local -a running_pids=()

    # do the job, restart updated services
    for service; do
        echo "systemctl restart $service"
        systemctl restart "$service" &
        if (( SERIALIZE )); then
            wait
            # display status directly when serialize and not quiet
            (( STATUS )) && systemctl --no-pager --lines=0 status "$service"
        else
            # register pids
            registered_pids[$!]="$service"
        fi
    done

    # display status as soon as available when not serialized
    while (( ${#registered_pids[*]} )); do
        # wait for at least one process to finish
        wait -n

        mapfile -t running_pids < <(jobs -p)

        # count registered pid for loop protection
        last_registered_pids_count=${#registered_pids[*]}

        for pid in "${!registered_pids[@]}"; do
            in_array "$pid" "${running_pids[@]}" && continue
            # show units status
            (( STATUS )) && systemctl --no-pager --lines=0 status "${registered_pids[$pid]}"
            unset "registered_pids[$pid]"
            break
        done

        # ensure we are not at 1st infinite loop
        # if we didn't remove a process something wrong happen
        if (( last_registered_pids_count == ${#registered_pids[*]} )); then
            error "Unable to wait processes to finish"
            error "Registered PIDs: ${registered_pids[*]}"
            error "Running PIDs: ${running_pids[*]}"
            break
        fi
    done
}

# reload or reexectute systemd
reload_systemd() {
  if grep -qF '(deleted)' /proc/1/maps; then
    arrow 'Restart systemd'
    systemctl --system daemon-reexec
  else
    arrow 'Reload systemd'
    systemctl --system daemon-reload
  fi
}

# display application usage and exit 2
usage() {
    echo "usage ${0##*/} [options]"
    echo "description: check for updated files in a service"
    echo "For toggle options: lowercase enables (1), uppercase disables (0)"
    echo 'options:'
    echo '  -h: this help' >&2
    echo "  -a/-A: auto confirmation (default: -A / $AUTOCONFIRM)" >&2
    echo "  -l/-L: call systemd daemon-(reload|reexec) (default: -l / $RELOAD)" >&2
    echo "  -f/-F: display failed services before quit (default: -f / $FAILED)" >&2
    echo "  -p/-P: call pacdiff before act (default: -p / $PACDIFF)" >&2
    echo "  -t/-T: view diffs in 3-way fashion with pacdiff (default: -T / $THREEWAY)" >&2
    echo "  -r/-R: restart services with updated files (default: -r / $RESTART)" >&2
    echo "  -s/-S: display status of restarted service (default: -s / $STATUS)" >&2
    echo "  -u/-U: act on services in users slice (default: -U / $USER_SLICE)" >&2
    echo "  -m/-M: act on services in machine slice (default: -M / $MACHINE_SLICE)" >&2
    echo "  -z/-Z: serialize action (default: -Z / $SERIALIZE)" >&2
    echo "  -i 'service_name'.service: ignore a specific service (can be used multiple times)" >&2
    exit 2
}

# parse command line arguments
# set options as global vars
argparse() {
    local opt
    while getopts 'AahFfLlPpRrSsTtUuMmZzi:' opt; do
        case $opt in
            A) AUTOCONFIRM=0;;      a) AUTOCONFIRM=1;;
            F) FAILED=0;;           f) FAILED=1;;
            L) RELOAD=0;;           l) RELOAD=1;;
            P) PACDIFF=0;;          p) PACDIFF=1;;
            R) RESTART=0;;          r) RESTART=1;;
            S) STATUS=0;;           s) STATUS=1;;
            T) THREEWAY=0;;         t) THREEWAY=1;;
            U) USER_SLICE=0;;       u) USER_SLICE=1;;
            M) MACHINE_SLICE=0;;    m) MACHINE_SLICE=1;;
            Z) SERIALIZE=0;;        z) SERIALIZE=1;;
            i) if [[ "$OPTARG" == *.service ]]; then
                    IGNORED_SERVICES+=("$OPTARG")
               else
                    usage
               fi
               ;;
            *) usage;;
        esac
    done
    shift $((OPTIND - 1));
    (( $# > 0 )) && usage
}

# emulated program entry point
main() {
    # avoid to be sighup'ed by interactive shell
    trap '' SIGHUP

    # parse command line options
    argparse "$@"

    # from now, we need to be root
    (( UID != 0 )) && error 'You need to be root' && exit 1

    # call pacdiff unless explicitly disabled to ensure config files are updated before restart
    if (( PACDIFF )); then
        if  command -v pacdiff &> /dev/null ; then
            arrow 'Run pacdiff'
            if (( THREEWAY )); then
                pacdiff --threeway
            else
                pacdiff
            fi
        else
            warn 'skipping pacdiff as not installed'
        fi
    fi

    # ensure systemd has been reloaded or reexectued
    (( RELOAD )) && reload_systemd

    arrow 'Services with broken maps files'
    local -a broken_services
    mapfile -t broken_services < <(get_broken_maps)
    echo "Found: ${#broken_services[@]}"
    if (( ${#broken_services[@]} )); then
        if (( RESTART )); then
            display_restart "${broken_services[@]}"
            if confirm 'Execute?'; then
                arrow 'Restart broken services'
                restart_services "${broken_services[@]}"
            fi
        else
             display_restart "${broken_services[@]}" | sed "s/systemctl restart //g"
        fi
    fi

    arrow 'Services missing on the system bus'
    local -a missing_services
    mapfile -t missing_services < <(get_missing_dbus)
    echo "Found: ${#missing_services[@]}"
    if (( ${#missing_services[@]} )); then
        if (( RESTART )); then
            display_restart "${missing_services[@]}"
            if confirm 'Execute?'; then
                arrow 'Restart missing services'
                restart_services "${missing_services[@]}"
            fi
        else
            display_restart "${missing_services[@]}" | sed "s/systemctl restart //g"
        fi
    fi

    # list only failed systemd units
    if (( FAILED )); then
        arrow "List failed units"
        systemctl --failed --all --no-pager --no-legend --full list-units
    fi
}

main "$@"

exit 0
