#!/bin/bash
# ===========================================================================
# train_bayes.sh - Bayes-Klassifikator von rspamd aus den vorhandenen
#                  Postfaechern anlernen.
#
# Hintergrund: rspamd traegt mit BAYES_* erst dann zum Score bei, wenn genug
# gelernt wurde (Default: 200 Spam UND 200 Ham). Bis dahin liefert der
# Klassifikator gar nichts - genau das war im Log zu sehen, dort tauchte kein
# einziges BAYES_-Symbol auf. Statt monatelang auf manuelles Verschieben in
# den Junk-Ordner zu warten (das macht RSPAMD_LEARN=1 bereits automatisch),
# lernt dieses Skript einmalig aus dem, was ohnehin auf der Platte liegt.
#
# AUSFUEHRUNG - immer IM Container:
#
#   docker exec mailserver bash /tmp/docker-mailserver/train_bayes.sh
#       -> Trockenlauf: zaehlt nur, lernt nichts. Das ist der Default.
#
#   docker exec mailserver bash /tmp/docker-mailserver/train_bayes.sh --apply
#       -> lernt tatsaechlich
#
# Optionen:
#   --apply             Wirklich lernen (ohne das passiert nichts)
#   --max N             Max. Nachrichten je Postfach und Klasse (Default 400)
#   --include-sent      Gesendete Mails zusaetzlich als Ham lernen
#   --only MUSTER       Nur Postfaecher, deren Adresse MUSTER enthaelt
#   --host HOST:PORT    rspamd-Controller (Default 127.0.0.1:11334)
#
# Passwort: Falls der Controller eines verlangt (bei euch in
# rspamd/override.d/worker-controller.inc gesetzt), vorweg setzen:
#   docker exec -e RSPAMC_PASSWORD='...' mailserver bash /tmp/.../train_bayes.sh
#
# WICHTIG zur Datenqualitaet:
#   - Ham kommt aus der INBOX. Liegt dort unerkannter Spam, wird er als Ham
#     gelernt und verschlechtert die Erkennung. Vorher grob durchsehen lohnt.
#   - Spam kommt aus Junk. Liegt dort faelschlich einsortierte legitime Mail,
#     gilt dasselbe umgekehrt.
#   - Ein sehr schiefes Verhaeltnis (z.B. 8000 Ham zu 50 Spam) macht den
#     Klassifikator einseitig. Der Trockenlauf zeigt das Verhaeltnis an -
#     danach ggf. mit --max nachsteuern.
# ===========================================================================

set -uo pipefail

MAILBASE="${MAILBASE:-/var/mail}"
RSPAMC_HOST="127.0.0.1:11334"
MAX_PER_CLASS=400
MAX_BYTES=5242880        # 5 MB - groessere Mails bringen fuers Lernen nichts
BATCH=50                 # Dateien pro rspamc-Aufruf
APPLY=0
INCLUDE_SENT=0
ONLY=""

usage() { sed -n '2,40p' "$0" | sed 's/^# \{0,1\}//' ; }

while [[ $# -gt 0 ]]; do
  case "$1" in
    --apply)        APPLY=1; shift ;;
    --max)          MAX_PER_CLASS="${2:-}"; shift 2 ;;
    --include-sent) INCLUDE_SENT=1; shift ;;
    --only)         ONLY="${2:-}"; shift 2 ;;
    --host)         RSPAMC_HOST="${2:-}"; shift 2 ;;
    -h|--help)      usage; exit 0 ;;
    *) echo "Unbekannte Option: $1" >&2; echo "Hilfe: $0 --help" >&2; exit 1 ;;
  esac
done

if [[ ! -d ${MAILBASE} ]]; then
  echo "FEHLER: ${MAILBASE} existiert nicht. Laeuft das Skript wirklich IM Container?" >&2
  exit 1
fi

RSPAMC=(rspamc -h "${RSPAMC_HOST}")
[[ -n ${RSPAMC_PASSWORD:-} ]] && RSPAMC+=(-P "${RSPAMC_PASSWORD}")

if ! command -v rspamc >/dev/null 2>&1; then
  echo "FEHLER: rspamc nicht gefunden." >&2
  exit 1
fi

if ! "${RSPAMC[@]}" stat >/dev/null 2>&1; then
  echo "FEHLER: Controller unter ${RSPAMC_HOST} nicht erreichbar oder Passwort noetig." >&2
  echo "        Versuche es mit:  docker exec -e RSPAMC_PASSWORD='...' mailserver bash $0 ..." >&2
  exit 1
fi

# ---------------------------------------------------------------------------
# Dateien einsammeln: bis LIMIT regulaere Dateien unterhalb der uebergebenen
# Maildir-Verzeichnisse, groessenbegrenzt. Neueste zuerst waere schoener, ist
# aber bei zehntausenden Dateien zu teuer - Maildir-Namen beginnen mit dem
# Timestamp, daher liefert die Sortierung eine brauchbare Naeherung.
# ---------------------------------------------------------------------------
gather() {
  local limit="$1"; shift
  local dir
  for dir in "$@"; do
    [[ -d ${dir} ]] || continue
    find "${dir}" -maxdepth 1 -type f -size -"${MAX_BYTES}"c -print 2>/dev/null
  done | sort -r | head -n "${limit}"
}

# ---------------------------------------------------------------------------
# Lernen in Batches. Gibt "gelernt fehlgeschlagen" zurueck.
# ---------------------------------------------------------------------------
learn() {
  local cmd="$1"; shift
  local -a files=("$@")
  local total=${#files[@]}
  local ok=0 fail=0 idx=0 out

  while (( idx < total )); do
    out=$("${RSPAMC[@]}" "${cmd}" "${files[@]:idx:BATCH}" 2>&1)
    ok=$(( ok + $(grep -c 'success = true' <<<"${out}") ))
    fail=$(( fail + $(grep -ci 'error' <<<"${out}") ))
    idx=$(( idx + BATCH ))
  done

  echo "${ok} ${fail}"
}

# ---------------------------------------------------------------------------
# Hauptlauf
# ---------------------------------------------------------------------------
if (( APPLY )); then
  echo "=== Bayes-Training (SCHREIBEND) ==="
else
  echo "=== Bayes-Training (Trockenlauf - es wird nichts gelernt) ==="
fi
echo "Basis: ${MAILBASE} | max ${MAX_PER_CLASS} je Postfach und Klasse"
[[ -n ${ONLY} ]] && echo "Filter: nur Postfaecher mit '${ONLY}'"
printf '\n%-42s %8s %8s\n' "POSTFACH" "SPAM" "HAM"
printf -- '---------------------------------------------------------------\n'

total_spam=0
total_ham=0
learned_spam=0
learned_ham=0
failed=0

for domain_dir in "${MAILBASE}"/*/; do
  [[ -d ${domain_dir} ]] || continue
  domain=$(basename "${domain_dir}")

  for user_dir in "${domain_dir}"*/; do
    user_dir="${user_dir%/}"
    [[ -d ${user_dir}/cur ]] || continue      # keine Maildir -> ueberspringen

    mailbox="$(basename "${user_dir}")@${domain}"
    [[ -n ${ONLY} && ${mailbox} != *"${ONLY}"* ]] && continue

    mapfile -t spam_files < <(gather "${MAX_PER_CLASS}" \
      "${user_dir}/.Junk/cur"  "${user_dir}/.Junk/new" \
      "${user_dir}/.Spam/cur"  "${user_dir}/.Spam/new")

    ham_dirs=("${user_dir}/cur" "${user_dir}/new")
    if (( INCLUDE_SENT )); then
      ham_dirs+=("${user_dir}/.Sent/cur" "${user_dir}/.Sent Messages/cur")
    fi
    mapfile -t ham_files < <(gather "${MAX_PER_CLASS}" "${ham_dirs[@]}")

    n_spam=${#spam_files[@]}
    n_ham=${#ham_files[@]}
    (( n_spam == 0 && n_ham == 0 )) && continue

    printf '%-42s %8d %8d\n' "${mailbox}" "${n_spam}" "${n_ham}"
    total_spam=$(( total_spam + n_spam ))
    total_ham=$(( total_ham + n_ham ))

    if (( APPLY )); then
      if (( n_spam > 0 )); then
        read -r ok fail <<<"$(learn learn_spam "${spam_files[@]}")"
        learned_spam=$(( learned_spam + ok )); failed=$(( failed + fail ))
      fi
      if (( n_ham > 0 )); then
        read -r ok fail <<<"$(learn learn_ham "${ham_files[@]}")"
        learned_ham=$(( learned_ham + ok )); failed=$(( failed + fail ))
      fi
    fi
  done
done

printf -- '---------------------------------------------------------------\n'
printf '%-42s %8d %8d\n\n' "SUMME" "${total_spam}" "${total_ham}"

if (( total_spam > 0 )); then
  ratio=$(( total_ham / (total_spam > 0 ? total_spam : 1) ))
  if (( ratio >= 5 )); then
    echo "WARNUNG: Ham/Spam-Verhaeltnis etwa ${ratio}:1 - sehr einseitig."
    echo "         Besser mit kleinerem --max fahren, z.B. --max $(( total_spam / 4 + 10 ))."
    echo
  fi
fi

if (( ! APPLY )); then
  echo "Trockenlauf beendet. Zum echten Lernen dasselbe Kommando mit --apply."
  echo "rspamd braucht je Klasse mindestens 200 Nachrichten, bevor BAYES_* feuert."
  exit 0
fi

echo "Gelernt: ${learned_spam} Spam, ${learned_ham} Ham, ${failed} Fehler/Duplikate."
echo
echo "=== rspamc stat ==="
"${RSPAMC[@]}" stat 2>/dev/null | grep -iE 'learn|classifier|statfile' || true
echo
echo "Kontrolle an der naechsten eingehenden Mail:"
echo "  docker exec mailserver sh -c 'grep rspamd_task_write_log /var/log/mail/rspamd.log | tail -3'"
echo "In der Symbolliste sollte jetzt BAYES_HAM oder BAYES_SPAM auftauchen."