#!/bin/bash

# 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

# disable grep options to avoid non default behaviour
unset GREP_OPTIONS

# Systemd cgroup path
SYSTEMD_CGROUP_BASE_PATH='/sys/fs/cgroup/systemd'

# colors
if [[ -t 1 ]]; then
    shopt -s xpg_echo
    c_arrow='\e[1;34m'
    c_title='\e[1;37m'
    c_svc='\e[1;35m'
    c_warn='\e[5;30;43m'
    c_error='\e[1;31m'
    c_rst='\e[m'
fi

# default options
autoconfirm=0       # autoconfirmation
dbus=1              # relauch when dbus
debug=0             # debug mode
failed=1            # display failed service at the end
pacdiff=1           # run pacdiff
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

# display application usage and exit 2
usage() {
    echo "usage ${0##*/} [options]"
    echo "description: check for updated files in a service"
    echo 'options:'
    echo '  -h: this help' >&2
    echo "  -d: debug mode" >&2
    echo "  -c: auto confirmation" >&2
    echo "  -b/-B: restart (or not) ${0##*/} if dbus was updated (default: $dbus)" >&2
    echo "  -l/-L: call (or not) systemd daemon-reload (default: $reload)" >&2
    echo "  -f/-F: display (or not) failed services before quit (default: $failed)" >&2
    echo "  -p/-P: call (or not) pacdiff before act (default: $pacdiff)" >&2
    echo "  -r/-R: restart (or not) services with updated files (default: $restart)" >&2
    echo "  -s/-S: display (or not) status of restarted service (default: $status)" >&2
    echo "  -u/-U: act (or not) on services in users slice (default: $user_slice)" >&2
    echo "  -z/-Z: serialize (or not) action (default: $serialize)" >&2
    exit 2
}

# print $* as an arrow line
arrow() {
    printf "%b==> %b%s%b\n" "$c_arrow" "$c_title" "$*" "$c_rst"
}

# print $* as an error message
error() {
    printf "%bError: %b%s%b\n" "$c_error" "$c_title" "$*" "$c_rst" >&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 in 5 4 3 2 1; 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
}

while getopts 'ahBbdFfLlPpRrSsUuZz' opt; do
    case $opt in
        a) autoconfirm=0;;
        B) dbus=0;;             b) dbus=1;;
        d) debug=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;;
        U) user_slice=0;;       u) user_slice=1;;
        Z) serialize=0;;        z) serialize=1;;
        *) usage;;
    esac
done
shift $((OPTIND - 1));
(( $# > 0 )) && usage

# avoid to be sighup'ed by interactive shell
trap '' SIGHUP

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

# call pacdiff
(( $pacdiff )) && arrow 'Run pacdiff' && pacdiff

# reload units list
(( $reload )) && arrow 'Reload systemd' && systemctl --system daemon-reload

# list of running services
arrow 'List runnings systemd services'
declare -a services
services=($(systemctl --no-legend --full --type service --state running|cut -f1 -d' '))

# list of bus names
arrow 'List Dbus clients'
declare -a buses
buses=($(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'))

# count beggar services
arrow "Search for updated mapped files"
declare -a needy=() pids=()
declare -i pid=0
for svc in "${services[@]}"; do
    unit_path="$(systemctl -p ControlGroup show "$svc"|cut -f 2 -d=)"
    busname="$(systemctl -p BusName show "$svc"|cut -f 2 -d=)"
    # get the right pidfile name
    unset 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 $svc." && continue
    # skip non system units
    (( $user_slice == 0 )) && [[ "$unit_path" =~ /user\.slice/ ]] && continue
    # parse pidfile
    pids=( $(< "$pidfile") )
    if (( "${#pids[*]}" == 0 )); then
        error "Unable to parse pid file for $svc."
        continue
    fi
    for pid in "${pids[@]}"; do
        maps_path="/proc/$pid/maps"
        [[ -r "$maps_path" ]] || {
            error "Unable to read maps file of $svc for pid $pid."
            continue
        }
        deleted=$(grep -F '(deleted)' "$maps_path" |sed -nr 's|^\S+ ..x. \S+ \S+ \S+ \s+||p'|sort|uniq)
        if [[ -n $deleted ]] || { [[ -n "$busname" ]] && ! in_array "$busname" "${buses[@]}"; }; then
            needy+=("$svc")
            if (( $debug )); then
                echo "${c_title}Service:${c_svc} $svc${c_rst}"
                echo "${c_title}Pid:${c_rst} $pid"
                echo "${c_title}Bus:${c_rst} $busname"
                echo -n "${c_title}BusName matching on the system bus: ${c_rst}"
                in_array "$busname" "${buses[@]}" && echo 'Yes' || echo 'No'
                echo "${c_title}Deleted files:${c_rst}"
                echo "$deleted"
                echo
            fi
            break
        fi
    done
done

# display what we will do
(( "${#needy[*]}" )) && echo '-------8<-------------------------------8<---------'
for svc in "${needy[@]}"; do
    echo "systemctl restart '$svc'"
done
(( "${#needy[*]}" )) && echo '-------8<-------------------------------8<---------'

# start the dangerous action below
if (( $restart == 1 && ${#needy[*]} > 0 )) && confirm 'Confirm service restart?'; then

    declare -A registered_pids=()
    declare -a running_pids=()
    declare -i last_registered_pids_count

    # do the job, restart updated services
    for svc in "${needy[@]}"; do
        echo "Restarting ${c_svc}$svc${c_rst}"
        systemctl restart "$svc" &
        if (( $serialize )); then
            wait
            # display status directly when serialize and not quiet
            (( $status )) && systemctl --no-pager --lines=0 status "$svc"
        else
            # register pids
            registered_pids[$!]="$svc"
        fi
    done

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

        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

    # warn if dbus was restart
    if in_array dbus.service "${needy[@]}"; then
        echo "${c_warn}After dbus restart, you should run ${0##*/} twice${c_rst}" >&2
        if (( $dbus )) && [[ -z "$CHECKSERVICE_WAS_RESTARTED" ]]; then
            echo "${c_warn}Doing it for you. No need to thanks me!${c_rst}" >&2
            export CHECKSERVICE_WAS_RESTARTED=1
            exec "$0" "$@"
        fi
    fi
fi

(( $failed )) && arrow "List failed units" && systemctl --failed --all --no-pager --no-legend --full list-units

exit 0