#!/usr/bin/env bash
#
# Continuously monitor the combined resident memory (RSS) of processes owned
# by one cPanel user. If the total exceeds the configured limit, terminate
# only that user's processes.
#
# Usage:
#   sudo ./elimsoft-memory-guard.sh --dry-run
#   sudo ./elimsoft-memory-guard.sh --enforce
#
# Optional environment overrides:
#   CPANEL_USER=elimsoft LIMIT_MB=50 INTERVAL_SECONDS=5 \
#     GRACE_SECONDS=3 COOLDOWN_SECONDS=10 \
#     LOCK_FILE=/run/lock/elimsoft-memory-guard.lock \
#     ./elimsoft-memory-guard.sh --enforce

set -Eeuo pipefail

readonly CPANEL_USER="${CPANEL_USER:-elimsoft}"
readonly LIMIT_MB="${LIMIT_MB:-50}"
readonly INTERVAL_SECONDS="${INTERVAL_SECONDS:-5}"
readonly GRACE_SECONDS="${GRACE_SECONDS:-3}"
readonly COOLDOWN_SECONDS="${COOLDOWN_SECONDS:-10}"
readonly PROGRAM_NAME="${0##*/}"
readonly LOCK_FILE="${LOCK_FILE:-/run/lock/elimsoft-memory-guard.lock}"

usage() {
  printf '%s\n' \
    "Usage: sudo $PROGRAM_NAME --dry-run|--enforce" \
    "" \
    "  --dry-run  Log threshold violations without killing processes." \
    "  --enforce  Send SIGTERM, wait, then SIGKILL any survivors." \
    "" \
    "Defaults: user=$CPANEL_USER limit=${LIMIT_MB}MB interval=${INTERVAL_SECONDS}s"
}

case "${1:-}" in
  --dry-run)
    readonly MODE="dry-run"
    ;;
  --enforce)
    readonly MODE="enforce"
    ;;
  --help|-h)
    usage
    exit 0
    ;;
  *)
    usage >&2
    exit 2
    ;;
esac

if (( EUID != 0 )); then
  printf 'ERROR: Run this script as root.\n' >&2
  exit 1
fi

for setting in LIMIT_MB INTERVAL_SECONDS GRACE_SECONDS COOLDOWN_SECONDS; do
  value="${!setting}"
  if [[ ! "$value" =~ ^[1-9][0-9]*$ ]]; then
    printf 'ERROR: %s must be a positive whole number; received %q.\n' \
      "$setting" "$value" >&2
    exit 1
  fi
done

if ! user_uid="$(id -u "$CPANEL_USER" 2>/dev/null)"; then
  printf 'ERROR: cPanel user %q does not exist.\n' "$CPANEL_USER" >&2
  exit 1
fi

if (( user_uid == 0 )); then
  printf 'ERROR: Refusing to monitor or terminate root-owned processes.\n' >&2
  exit 1
fi

for required_command in awk date flock kill ps sleep; do
  if ! command -v "$required_command" >/dev/null 2>&1; then
    printf 'ERROR: Required command %q was not found.\n' \
      "$required_command" >&2
    exit 1
  fi
done

if ! { exec 9>"$LOCK_FILE"; } 2>/dev/null; then
  printf 'ERROR: Cannot open lock file %q.\n' "$LOCK_FILE" >&2
  exit 1
fi
if ! flock -n 9; then
  printf 'ERROR: Another instance is already running.\n' >&2
  exit 1
fi

readonly USER_UID="$user_uid"
readonly LIMIT_KB=$((LIMIT_MB * 1024))

log_message() {
  local message="$1"
  printf '%s %s\n' "$(date '+%F %T')" "$message"

  if command -v logger >/dev/null 2>&1; then
    logger -t "$PROGRAM_NAME" -- "$message"
  fi
}

total_rss_kb() {
  { ps -u "$USER_UID" -o rss= 2>/dev/null || true; } |
    awk '{ total += $1 } END { printf "%d\n", total + 0 }'
}

get_user_pids() {
  { ps -u "$USER_UID" -o pid= 2>/dev/null || true; } |
    awk '$1 ~ /^[0-9]+$/ { print $1 }'
}

process_summary() {
  { ps -u "$USER_UID" -o pid=,rss=,comm= --sort=-rss 2>/dev/null || true; } |
    awk 'NR <= 10 {
      printf "%sPID=%s RSS=%sKB CMD=%s", separator, $1, $2, $3
      separator = "; "
    }'
}

terminate_user_processes() {
  local -a pids=()
  local -a survivors=()

  mapfile -t pids < <(get_user_pids)
  if (( ${#pids[@]} == 0 )); then
    log_message "No processes remained for user $CPANEL_USER."
    return
  fi

  log_message "Sending SIGTERM to ${#pids[@]} process(es) owned by $CPANEL_USER."
  kill -TERM "${pids[@]}" 2>/dev/null || true
  sleep "$GRACE_SECONDS"

  mapfile -t survivors < <(get_user_pids)
  if (( ${#survivors[@]} > 0 )); then
    log_message "Sending SIGKILL to ${#survivors[@]} surviving process(es) owned by $CPANEL_USER."
    kill -KILL "${survivors[@]}" 2>/dev/null || true
  fi
}

stop_requested=0
handle_stop() {
  stop_requested=1
  log_message "Stop requested; exiting monitor."
}
trap handle_stop INT TERM

log_message "Started mode=$MODE user=$CPANEL_USER uid=$USER_UID limit=${LIMIT_MB}MB interval=${INTERVAL_SECONDS}s."

while (( stop_requested == 0 )); do
  rss_kb="$(total_rss_kb)"

  if (( rss_kb > LIMIT_KB )); then
    rss_mb="$(awk -v kb="$rss_kb" 'BEGIN { printf "%.1f", kb / 1024 }')"
    summary="$(process_summary)"
    log_message "LIMIT EXCEEDED: user=$CPANEL_USER combined_RSS=${rss_mb}MB limit=${LIMIT_MB}MB. $summary"

    if [[ "$MODE" == "enforce" ]]; then
      terminate_user_processes
    else
      log_message "Dry run: no processes were terminated."
    fi

    sleep "$COOLDOWN_SECONDS" &
  else
    sleep "$INTERVAL_SECONDS" &
  fi

  wait $! || true
done

exit 0