#!/usr/bin/env bash
# ------------------------------------------------------------------------------
# Copyright (c) 2013-2022 Geoffray Levasseur <fatalerrors@geoffray-levasseur.org>
# Protected by the BSD3 license. Please read bellow for details.
#
# * Redistribution and use in source and binary forms,
# * with or without modification, are permitted provided
# * that the following conditions are met:
# *
# *   Redistributions of source code must retain the above
# *   copyright notice, this list of conditions and the
# *   following disclaimer.
# *
# *   Redistributions in binary form must reproduce the above
# *   copyright notice, this list of conditions and the following
# *   disclaimer in the documentation and/or other materials
# *   provided with the distribution.
# *
# *   Neither the name of the copyright holder nor the names
# *   of any other contributors may be used to endorse or
# *   promote products derived from this software without
# *   specific prior written permission.
# *
# * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
# * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
# * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
# * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# * OF SUCH DAMAGE.
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
# Search processes matching the given string
# Usage: ppg <string>
ppg()
{
    local PARSED

    PARSED=$(getopt -o h --long help -n 'ppg' -- "$@")
    # shellcheck disable=SC2181 # getopt return code is checked immediately after
    if [[ $? -ne 0 ]]; then
        disp E "Invalid options, use \"ppg --help\" to display usage."
        return 1
    fi

    eval set -- "$PARSED"
    while true; do
        case "$1" in
            -h|--help)
                printf "ppg: Search processes matching the given string.\n\n"
                printf "Usage: ppg <string>\n\n"
                printf "Options:\n"
                printf "\t-h, --help\t\tDisplay this help screen\n"
                return 0
                ;;
            --)
                shift
                break
                ;;
            *)
                disp E "Invalid options, use \"ppg --help\" to display usage."
                return 1
                ;;
        esac
    done

    if [[ -z "$1" ]]; then
        disp E "Usage: ppg <string>"
        return 1
    fi

    local pattern="$*"

    if command -v pgrep >/dev/null 2>&1; then
        pgrep -af -- "$pattern"
        return $?
    fi

    ps -ef | awk -v pattern="$pattern" '
        NR == 1 {
            print
            next
        }
        index($0, pattern) {
            print
            matched = 1
        }
        END {
            exit matched ? 0 : 1
        }
    '
}
export -f ppg
# ------------------------------------------------------------------------------


# ------------------------------------------------------------------------------
# List processes owned by a specific user
# Usage: ppu <username>
ppu()
{
    local PARSED

    PARSED=$(getopt -o h --long help -n 'ppu' -- "$@")
    # shellcheck disable=SC2181 # getopt return code is checked immediately after
    if [[ $? -ne 0 ]]; then
        disp E "Invalid options, use \"ppu --help\" to display usage."
        return 1
    fi

    eval set -- "$PARSED"
    while true; do
        case "$1" in
            -h|--help)
                printf "ppu: List processes owned by a specific user.\n\n"
                printf "Usage: ppu <username>\n\n"
                printf "Options:\n"
                printf "\t-h, --help\t\tDisplay this help screen\n"
                return 0
                ;;
            --)
                shift
                break
                ;;
            *)
                disp E "Invalid options, use \"ppu --help\" to display usage."
                return 1
                ;;
        esac
    done

    if [[ -z "$1" ]]; then
        disp E "Usage: ppu <username>"
        return 1
    fi

    # -u lists processes for a specific user
    # -o provides a clean, standard output format
    ps -u "$1" -o "${PPU_DEFAULT_FORMAT:-pid,user,%cpu,%mem,start,time,command}"
}
export -f ppu
# ------------------------------------------------------------------------------


# ------------------------------------------------------------------------------
# List processes by exact command name (no path/parameters)
# Usage: ppn <command_name>
ppn()
{
    local PARSED

    PARSED=$(getopt -o h --long help -n 'ppn' -- "$@")
    # shellcheck disable=SC2181 # getopt return code is checked immediately after
    if [[ $? -ne 0 ]]; then
        disp E "Invalid options, use \"ppn --help\" to display usage."
        return 1
    fi

    eval set -- "$PARSED"
    while true; do
        case "$1" in
            -h|--help)
                printf "ppn: List processes by exact command name (no path/parameters).\n\n"
                printf "Usage: ppn <command_name>\n\n"
                printf "Options:\n"
                printf "\t-h, --help\t\tDisplay this help screen\n"
                return 0
                ;;
            --)
                shift
                break
                ;;
            *)
                disp E "Invalid options, use \"ppn --help\" to display usage."
                return 1
                ;;
        esac
    done

    if [[ -z "$1" ]]; then
        disp E "Usage: ppn <command_name>"
        return 1
    fi

    # -e: select all processes
    # -o: specify custom output columns (PID and Command name)
    # grep -w: ensures exact word matching so 'bash' doesn't match 'dbash'
    # shellcheck disable=SC2009 # pgrep do not offer the -w switch
    ps -eo pid,comm | grep -w "$1"
}
export -f ppn
# ------------------------------------------------------------------------------


# ------------------------------------------------------------------------------
# Get PID list of the given process name
# Usage: gpid <process_name [process_name2 ...]>
gpid()
{
    local PARSED

    PARSED=$(getopt -o h --long help -n 'gpid' -- "$@")
    # shellcheck disable=SC2181 # getopt return code is checked immediately after
    if [[ $? -ne 0 ]]; then
        disp E "Invalid options, use \"gpid --help\" to display usage."
        return 1
    fi

    eval set -- "$PARSED"
    while true; do
        case "$1" in
            -h|--help)
                printf "gpid: Get PID list of the given process name.\n\n"
                printf "Usage: gpid <process_name [process_name2 ...]>\n\n"
                printf "Options:\n"
                printf "\t-h, --help\t\tDisplay this help screen\n"
                return 0
                ;;
            --)
                shift
                break
                ;;
            *)
                disp E "Invalid options, use \"gpid --help\" to display usage."
                return 1
                ;;
        esac
    done

    if [[ -z "$1" ]]; then
        disp E "Usage: gpid <process_name [process_name2 ...]>"
        return 1
    fi
    local single=0
    local found=0
    local proc_name result

    [[ $# -eq 1 ]] && single=1

    for proc_name in "$@"; do
        result=""

        if command -v pgrep >/dev/null 2>&1; then
            result=$(pgrep -d ' ' -x -- "$proc_name")
        else
            result=$(ps -eo pid=,comm= | awk -v proc="$proc_name" '
                $2 == proc {
                    if (out != "") {
                        out = out " "
                    }
                    out = out $1
                }
                END {
                    print out
                }
            ')
        fi

        [[ -z "$result" ]] && continue

        found=1
        if (( single )); then
            echo "$result"
        else
            echo "$proc_name: $result"
        fi
    done

    (( found )) || return 1
}
export -f gpid
# ------------------------------------------------------------------------------


# ------------------------------------------------------------------------------
# Kill all processes owned by the given users (kill user)
# Usage: ku <username1 [username2 ...]>
ku()
{
    local PARSED
    local dry_run=0
    local -a signal_opt=()

    PARSED=$(getopt -o hns: --long help,dry-run,signal: -n 'ku' -- "$@")
    # shellcheck disable=SC2181 # getopt return code is checked immediately after
    if [[ $? -ne 0 ]]; then
        disp E "Invalid options, use \"ku --help\" to display usage."
        return 1
    fi

    eval set -- "$PARSED"
    while true; do
        case "$1" in
            -h|--help)
                printf "ku: Kill all processes owned by the given users.\n\n"
                printf "Usage: ku [options] <username1 [username2 ...]>\n\n"
                printf "Options:\n"
                printf "\t-h, --help\t\tDisplay this help screen\n"
                printf "\t-n, --dry-run\t\tDisplay commands without executing them\n"
                printf "\t-s, --signal SIG\tSignal to send (overrides KU_DEFAULT_SIGNAL)\n"
                printf "\t    --signal=SIG\tSame as above\n"
                printf "\t    -SIG / -NUM\t\tSignal format compatible with kill\n"
                return 0
                ;;
            -n|--dry-run)
                dry_run=1
                shift
                ;;
            -s|--signal)
                signal_opt=(-s "$2")
                shift 2
                ;;
            --)
                shift
                break
                ;;
            -*)
                signal_opt=("$1")
                shift
                ;;
            *)
                break
                ;;
        esac
    done

    # Accept kill-style signal forms not handled by getopt, before usernames.
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -[0-9]*|-SIG*|-[[:alpha:]]*)
                signal_opt=("$1")
                shift
                ;;
            --)
                shift
                break
                ;;
            -*)
                disp E "Unknown option: $1, use \"ku --help\" to display usage."
                return 1
                ;;
            *)
                break
                ;;
        esac
    done

    if [[ -z "$1" ]]; then
        disp E "Usage: ku [options] <username1 [username2 ...]>"
        return 1
    fi

    local u
    for u in "$@"; do
        if ! id "$u" >/dev/null 2>&1; then
            disp E "User '$u' does not exist."
            return 1
        else
            local cmd

            # killall (psmisc) preferred; fall back to pkill (procps-ng).
            if command -v killall >/dev/null 2>&1; then
                cmd=(killall)
                if [[ ${#signal_opt[@]} -gt 0 ]]; then
                    cmd+=("${signal_opt[@]}")
                elif [[ -n "${KU_DEFAULT_SIGNAL:-}" ]]; then
                    cmd+=("-${KU_DEFAULT_SIGNAL}")
                fi
                cmd+=(-u "$u")
            elif command -v pkill >/dev/null 2>&1; then
                cmd=(pkill)
                # Translate killall's -s SIGNAME form to pkill's -SIGNAME form.
                if [[ ${#signal_opt[@]} -eq 2 && "${signal_opt[0]}" == "-s" ]]; then
                    cmd+=("-${signal_opt[1]}")
                elif [[ ${#signal_opt[@]} -gt 0 ]]; then
                    cmd+=("${signal_opt[@]}")
                elif [[ -n "${KU_DEFAULT_SIGNAL:-}" ]]; then
                    cmd+=("-${KU_DEFAULT_SIGNAL}")
                fi
                cmd+=(-u "$u")
            else
                disp E "ku: neither 'killall' (psmisc) nor 'pkill' (procps) is available."
                return 1
            fi

            if (( dry_run )); then
                printf "DRY-RUN: "
                printf "%q " "${cmd[@]}"
                printf "\n"
            else
                "${cmd[@]}"
            fi
        fi
    done
}
export -f ku
# ------------------------------------------------------------------------------


# ------------------------------------------------------------------------------
# Kill all children of a process then the process (kill tree)
# Usage: kt <pid> [kill_options]
kt()
{
    local PARSED
    local dry_run=0
    local -a pre_kill_opts=()

    PARSED=$(getopt -o hns: --long help,dry-run,signal: -n 'kt' -- "$@")
    # shellcheck disable=SC2181 # getopt return code is checked immediately after
    if [[ $? -ne 0 ]]; then
        disp E "Invalid options, use \"kt --help\" to display usage."
        return 1
    fi

    eval set -- "$PARSED"
    while true; do
        case "$1" in
            -h|--help)
                printf "kt: Kill all children of a process then the process (kill tree).\n\n"
                printf "Usage: kt [options] <pid> [kill_options]\n\n"
                printf "Options:\n"
                printf "\t-h, --help\t\tDisplay this help screen\n"
                printf "\t-n, --dry-run\t\tDisplay kill commands without executing them\n"
                printf "\t-s, --signal SIG\tSignal to send to process tree\n"
                printf "\t    --signal=SIG\tSame as above\n"
                printf "\t    -SIG / -NUM\t\tSignal format compatible with kill\n"
                return 0
                ;;
            -n|--dry-run)
                dry_run=1
                shift
                ;;
            -s|--signal)
                pre_kill_opts+=(-s "$2")
                shift 2
                ;;
            --)
                shift
                break
                ;;
            -*)
                pre_kill_opts+=("$1")
                shift
                ;;
            *)
                break
                ;;
        esac
    done

    # Accept kill-style signal forms not handled by getopt, before the PID.
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -[0-9]*|-SIG*|-[[:alpha:]]*)
                pre_kill_opts+=("$1")
                shift
                ;;
            --)
                shift
                break
                ;;
            -*)
                disp E "Unknown option: $1, use \"kt --help\" to display usage."
                return 1
                ;;
            *)
                break
                ;;
        esac
    done

    if [[ -z "$1" ]]; then
        disp E "Usage: kt [options] <pid> [kill_options]"
        return 1
    fi

    local parent_pid="$1"
    shift

    local -a kill_opts=("${pre_kill_opts[@]}" "$@")

    if [[ "$parent_pid" == "0" || "$parent_pid" == "1" ]]; then
        disp E "Safety abort: Refusing to kill PID $parent_pid (system critical)."
        return 1
    fi

    local children_pids
    children_pids=$(pgrep -P "$parent_pid" 2>/dev/null || true)

    local pid
    for pid in $children_pids; do
        if (( dry_run )); then
            kt --dry-run "$pid" "${kill_opts[@]}" || break
        else
            kt "$pid" "${kill_opts[@]}" || break
        fi
    done

    if (( dry_run )); then
        local cmd=(kill "${kill_opts[@]}" "$parent_pid")
        printf "DRY-RUN: "
        printf "%q " "${cmd[@]}"
        printf "\n"
    else
        kill "${kill_opts[@]}" "$parent_pid"
    fi
}
export -f kt
# ------------------------------------------------------------------------------


load_conf "processes"

# EOF