Initial commit for the rspamd spam/ham trainer

This commit is contained in:
2026-06-22 16:17:38 +02:00
commit 7bc9d7a901
18 changed files with 1071 additions and 0 deletions
+3
View File
@@ -0,0 +1,3 @@
"""Rspamd Mail Trainer package."""
__version__ = "0.1.0"
+3
View File
@@ -0,0 +1,3 @@
from .cli import main
raise SystemExit(main())
+619
View File
@@ -0,0 +1,619 @@
#!/usr/bin/env python3
"""
rspamd-mail-trainer.py
Find Maildir or mbox messages by sender/recipient/subject/date and train Rspamd as spam or ham.
Examples:
# Show likely mail roots for user johannes
./rspamd-mail-trainer.py roots --user johannes
# Dry-run: list matching messages from a sender
./rspamd-mail-trainer.py search --user johannes --from no-reply@ezweb.ne.jp
# Train matching messages as spam, asking for confirmation
./rspamd-mail-trainer.py train spam --user johannes --from no-reply@ezweb.ne.jp
# Train matching messages as ham by subject
./rspamd-mail-trainer.py train ham --user johannes --subject "Rechnung"
# Analyse matching messages with rspamc, but do not train
./rspamd-mail-trainer.py check --user johannes --from no-reply@ezweb.ne.jp
"""
from __future__ import annotations
import argparse
import datetime as dt
import email
import email.policy
import getpass
import hashlib
import mailbox
import os
import pwd
import re
import shutil
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from email.header import decode_header, make_header
from email.message import Message
from pathlib import Path
from typing import Iterable, Iterator, Optional
DEFAULT_MAX_SIZE = 25 * 1024 * 1024
MBOX_NAMES = {"mbox", "johannes"} # used only as a weak hint; actual detection is content-based
@dataclass(frozen=True)
class Candidate:
source_type: str # maildir-file or mbox-message
path: Path
message_index: Optional[int] # only for mbox
digest: str
size: int
sender: str
to: str
subject: str
date_header: str
@dataclass(frozen=True)
class SearchOptions:
from_re: Optional[re.Pattern[str]]
to_re: Optional[re.Pattern[str]]
subject_re: Optional[re.Pattern[str]]
any_re: Optional[re.Pattern[str]]
after: Optional[dt.datetime]
before: Optional[dt.datetime]
max_size: int
include_seen_duplicates: bool
def eprint(*args: object) -> None:
print(*args, file=sys.stderr)
def run(cmd: list[str], check: bool = False) -> subprocess.CompletedProcess[str]:
return subprocess.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=check)
def decode_header_value(value: object) -> str:
if value is None:
return ""
try:
return str(make_header(decode_header(str(value))))
except Exception:
return str(value)
def parse_message_bytes(raw: bytes) -> Message:
return email.message_from_bytes(raw, policy=email.policy.default)
def get_text_for_any_match(msg: Message, max_chars: int = 200_000) -> str:
parts: list[str] = []
for key in ("From", "To", "Cc", "Bcc", "Subject", "Date", "Message-ID", "Return-Path"):
parts.append(f"{key}: {decode_header_value(msg.get(key))}")
if msg.is_multipart():
for part in msg.walk():
ctype = part.get_content_type()
disp = (part.get_content_disposition() or "").lower()
if disp == "attachment":
continue
if ctype in {"text/plain", "text/html"}:
try:
parts.append(part.get_content())
except Exception:
pass
if sum(len(p) for p in parts) >= max_chars:
break
else:
try:
if msg.get_content_type() in {"text/plain", "text/html"}:
parts.append(msg.get_content())
except Exception:
pass
return "\n".join(parts)[:max_chars]
def parse_date_header(value: str) -> Optional[dt.datetime]:
if not value:
return None
try:
parsed = email.utils.parsedate_to_datetime(value)
if parsed is None:
return None
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=dt.timezone.utc)
return parsed.astimezone(dt.timezone.utc)
except Exception:
return None
def parse_user_date(value: Optional[str], end_of_day: bool = False) -> Optional[dt.datetime]:
if not value:
return None
try:
if re.fullmatch(r"\d{4}-\d{2}-\d{2}", value):
d = dt.date.fromisoformat(value)
t = dt.time(23, 59, 59) if end_of_day else dt.time(0, 0, 0)
return dt.datetime.combine(d, t, tzinfo=dt.timezone.utc)
# Accept ISO strings with timezone or without timezone.
parsed = dt.datetime.fromisoformat(value.replace("Z", "+00:00"))
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=dt.timezone.utc)
return parsed.astimezone(dt.timezone.utc)
except Exception as exc:
raise argparse.ArgumentTypeError(f"Invalid date '{value}': {exc}") from exc
def maybe_compile(pattern: Optional[str], ignore_case: bool, literal: bool) -> Optional[re.Pattern[str]]:
if not pattern:
return None
flags = re.IGNORECASE if ignore_case else 0
if literal:
pattern = re.escape(pattern)
try:
return re.compile(pattern, flags)
except re.error as exc:
raise argparse.ArgumentTypeError(f"Invalid regex '{pattern}': {exc}") from exc
def message_digest(raw: bytes) -> str:
return hashlib.sha256(raw).hexdigest()
def make_candidate(
*,
source_type: str,
path: Path,
raw: bytes,
msg: Message,
message_index: Optional[int] = None,
) -> Candidate:
return Candidate(
source_type=source_type,
path=path,
message_index=message_index,
digest=message_digest(raw),
size=len(raw),
sender=decode_header_value(msg.get("From")),
to=decode_header_value(msg.get("To")),
subject=decode_header_value(msg.get("Subject")),
date_header=decode_header_value(msg.get("Date")),
)
def candidate_matches(candidate: Candidate, msg: Message, opts: SearchOptions) -> bool:
if candidate.size > opts.max_size:
return False
msg_date = parse_date_header(candidate.date_header)
if opts.after and msg_date and msg_date < opts.after:
return False
if opts.before and msg_date and msg_date > opts.before:
return False
if opts.from_re and not opts.from_re.search(candidate.sender):
return False
if opts.to_re and not opts.to_re.search(candidate.to):
return False
if opts.subject_re and not opts.subject_re.search(candidate.subject):
return False
if opts.any_re:
text = get_text_for_any_match(msg)
if not opts.any_re.search(text):
return False
return True
def is_probably_mail_file(path: Path) -> bool:
# Maildir files may contain colon/comma flags and no extension. Avoid obvious non-mail files.
name = path.name
if name.startswith(".") and name not in {".mh_sequences"}:
return False
if path.suffix.lower() in {".lock", ".log", ".json", ".sqlite", ".db", ".index", ".cache", ".tmp"}:
return False
return True
def read_head(path: Path, n: int = 4096) -> bytes:
try:
with path.open("rb") as f:
return f.read(n)
except Exception:
return b""
def looks_like_mbox(path: Path) -> bool:
head = read_head(path, 4096)
return head.startswith(b"From ") and b"\nFrom:" in head[:4096]
def looks_like_single_message(path: Path) -> bool:
head = read_head(path, 8192)
# Common full-message headers. Do not require From: because spam may be malformed.
header_hits = sum(
token in head
for token in (b"\nFrom:", b"\nTo:", b"\nSubject:", b"\nDate:", b"\nMessage-ID:", b"\nReturn-Path:")
)
return header_hits >= 2
def iter_maildir_files(root: Path, max_file_size: int) -> Iterator[Path]:
for dirpath, dirnames, filenames in os.walk(root, followlinks=False):
# prune noisy dirs
dirnames[:] = [
d for d in dirnames
if d not in {".git", "tmp", "dovecot.index.cache"} and not d.endswith(".lock")
]
p = Path(dirpath)
# Be permissive: direct files under cur/new and files in Maildir-ish folders.
is_mail_leaf = p.name in {"cur", "new"} or "Maildir" in p.parts
if not is_mail_leaf:
continue
for filename in filenames:
f = p / filename
try:
st = f.stat()
except OSError:
continue
if not st.st_size or st.st_size > max_file_size:
continue
if is_probably_mail_file(f):
yield f
def iter_single_file_candidates(path: Path, opts: SearchOptions) -> Iterator[tuple[Candidate, bytes]]:
try:
raw = path.read_bytes()
msg = parse_message_bytes(raw)
except Exception as exc:
eprint(f"Skip unreadable mail file: {path}: {exc}")
return
cand = make_candidate(source_type="maildir-file", path=path, raw=raw, msg=msg)
if candidate_matches(cand, msg, opts):
yield cand, raw
def iter_mbox_candidates(path: Path, opts: SearchOptions) -> Iterator[tuple[Candidate, bytes]]:
try:
mbox = mailbox.mbox(str(path), create=False)
except Exception as exc:
eprint(f"Skip unreadable mbox: {path}: {exc}")
return
for index, _key in enumerate(mbox.keys()):
try:
m = mbox.get_message(_key)
raw = bytes(m)
msg = parse_message_bytes(raw)
except Exception as exc:
eprint(f"Skip unreadable mbox message {path}#{index}: {exc}")
continue
cand = make_candidate(source_type="mbox-message", path=path, raw=raw, msg=msg, message_index=index)
if candidate_matches(cand, msg, opts):
yield cand, raw
def discover_roots(user: Optional[str]) -> list[Path]:
roots: list[Path] = []
user = user or getpass.getuser()
# Dovecot mail_location if available.
doveconf = shutil.which("doveconf")
if doveconf:
cp = run([doveconf, "-n"])
if cp.returncode == 0:
for line in cp.stdout.splitlines():
if "mail_location" not in line:
continue
_, _, loc = line.partition("=")
loc = loc.strip()
loc = loc.split(":", 1)[1] if ":" in loc else loc
# Extract path prefix from common patterns.
loc = loc.split(":", 1)[0]
loc = loc.replace("%u", user)
if "%n" in loc or "%d" in loc:
# Keep only the static root before templates.
static = loc.split("%", 1)[0].rstrip("/")
if static:
roots.append(Path(static))
elif loc.startswith("~/"):
try:
home = Path(pwd.getpwnam(user).pw_dir)
except KeyError:
home = Path.home()
roots.append(home / loc[2:])
elif loc:
roots.append(Path(loc))
# Common locations.
try:
home = Path(pwd.getpwnam(user).pw_dir)
except KeyError:
home = Path(f"/home/{user}")
roots.extend([
home / "Maildir",
home / "mail",
home / ".mail",
home,
Path("/var/vmail"),
Path("/var/mail") / user,
Path("/var/spool/mail") / user,
])
# Deduplicate while preserving order, keep existing readable-ish paths.
seen: set[str] = set()
result: list[Path] = []
for r in roots:
try:
rp = r.expanduser().resolve(strict=False)
except Exception:
rp = r.expanduser()
key = str(rp)
if key in seen:
continue
seen.add(key)
if rp.exists():
result.append(rp)
return result
def iter_candidates(roots: list[Path], opts: SearchOptions) -> Iterator[tuple[Candidate, bytes]]:
seen_digests: set[str] = set()
for root in roots:
if not root.exists():
eprint(f"Skip missing root: {root}")
continue
if root.is_file():
paths = [root]
else:
paths = list(iter_maildir_files(root, opts.max_size))
# Also detect mbox files directly below selected root and a few common mail dirs.
for dirpath, dirnames, filenames in os.walk(root, followlinks=False):
depth = len(Path(dirpath).relative_to(root).parts) if Path(dirpath) != root else 0
if depth > 3:
dirnames[:] = []
continue
for filename in filenames:
f = Path(dirpath) / filename
if f in paths:
continue
try:
st = f.stat()
except OSError:
continue
if not st.st_size or st.st_size > opts.max_size * 200:
continue
if looks_like_mbox(f):
paths.append(f)
for p in paths:
if not p.exists() or not p.is_file():
continue
try:
if looks_like_mbox(p):
iterator = iter_mbox_candidates(p, opts)
elif looks_like_single_message(p):
iterator = iter_single_file_candidates(p, opts)
else:
continue
for cand, raw in iterator:
if not opts.include_seen_duplicates and cand.digest in seen_digests:
continue
seen_digests.add(cand.digest)
yield cand, raw
except PermissionError:
eprint(f"Permission denied: {p}")
except OSError as exc:
eprint(f"Cannot read {p}: {exc}")
def format_candidate(c: Candidate, idx: int) -> str:
where = str(c.path)
if c.message_index is not None:
where += f"#{c.message_index}"
return (
f"[{idx}] {c.source_type} size={c.size} sha256={c.digest[:12]}\n"
f" path: {where}\n"
f" from: {c.sender}\n"
f" to: {c.to}\n"
f" subject: {c.subject}\n"
f" date: {c.date_header}"
)
def create_temp_message(raw: bytes) -> Path:
tmp = tempfile.NamedTemporaryFile(prefix="rspamd-train-", suffix=".eml", delete=False)
with tmp:
tmp.write(raw)
return Path(tmp.name)
def call_rspamc(action: str, raw: bytes, path_hint: Path, rspamc: str, controller: Optional[str]) -> tuple[int, str, str]:
# For Maildir single-file messages we can pass the file directly. For mbox messages,
# raw bytes must be extracted to a temp .eml.
temp_path: Optional[Path] = None
target = path_hint
try:
if not path_hint.exists() or looks_like_mbox(path_hint):
temp_path = create_temp_message(raw)
target = temp_path
cmd = [rspamc]
if controller:
cmd.extend(["-h", controller])
cmd.extend([action, str(target)])
cp = run(cmd)
return cp.returncode, cp.stdout, cp.stderr
finally:
if temp_path:
try:
temp_path.unlink()
except OSError:
pass
def confirm(prompt: str) -> bool:
answer = input(f"{prompt} [y/N] ").strip().lower()
return answer in {"y", "yes", "j", "ja"}
def build_options(args: argparse.Namespace) -> SearchOptions:
return SearchOptions(
from_re=maybe_compile(args.from_, args.ignore_case, args.literal),
to_re=maybe_compile(args.to, args.ignore_case, args.literal),
subject_re=maybe_compile(args.subject, args.ignore_case, args.literal),
any_re=maybe_compile(args.any, args.ignore_case, args.literal),
after=parse_user_date(args.after, end_of_day=False),
before=parse_user_date(args.before, end_of_day=True),
max_size=args.max_size,
include_seen_duplicates=args.include_duplicates,
)
def resolve_roots(args: argparse.Namespace) -> list[Path]:
roots = [Path(p) for p in (args.root or [])]
if not roots:
roots = discover_roots(args.user)
return roots
def add_common_filters(parser: argparse.ArgumentParser) -> None:
parser.add_argument("--user", help="Unix user whose mailbox should be searched, e.g. johannes")
parser.add_argument("--root", action="append", help="Mail root/file to search. Can be used multiple times.")
parser.add_argument("--from", dest="from_", help="Regex or literal match against the From header")
parser.add_argument("--to", help="Regex or literal match against the To header")
parser.add_argument("--subject", help="Regex or literal match against the Subject header")
parser.add_argument("--any", help="Regex or literal match against selected headers and text body")
parser.add_argument("--after", help="Only messages after this date, e.g. 2026-06-01 or ISO timestamp")
parser.add_argument("--before", help="Only messages before this date, e.g. 2026-06-13 or ISO timestamp")
parser.add_argument("--literal", action="store_true", help="Treat --from/--to/--subject/--any as literal text, not regex")
parser.add_argument("--case-sensitive", dest="ignore_case", action="store_false", help="Use case-sensitive matching")
parser.set_defaults(ignore_case=True)
parser.add_argument("--max-size", type=int, default=DEFAULT_MAX_SIZE, help=f"Max single message size in bytes, default {DEFAULT_MAX_SIZE}")
parser.add_argument("--include-duplicates", action="store_true", help="Do not suppress duplicate messages by SHA-256 digest")
parser.add_argument("--limit", type=int, default=50, help="Maximum matches to process/display, default 50")
def main(argv: Optional[list[str]] = None) -> int:
parser = argparse.ArgumentParser(
description="Find Maildir/mbox messages and train Rspamd as spam or ham.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
sub = parser.add_subparsers(dest="command", required=True)
p_roots = sub.add_parser("roots", help="Show discovered mail roots")
p_roots.add_argument("--user", help="Unix user, e.g. johannes")
p_search = sub.add_parser("search", help="Search matching messages; does not train")
add_common_filters(p_search)
p_check = sub.add_parser("check", help="Run rspamc against matching messages; does not train")
add_common_filters(p_check)
p_check.add_argument("--rspamc", default=shutil.which("rspamc") or "rspamc", help="Path to rspamc")
p_check.add_argument("--controller", help="Rspamd controller address, e.g. 127.0.0.1:11334")
p_train = sub.add_parser("train", help="Train matching messages as spam or ham")
p_train.add_argument("class_", choices=["spam", "ham"], help="Training class")
add_common_filters(p_train)
p_train.add_argument("--rspamc", default=shutil.which("rspamc") or "rspamc", help="Path to rspamc")
p_train.add_argument("--controller", help="Rspamd controller address, e.g. 127.0.0.1:11334")
p_train.add_argument("--yes", "-y", action="store_true", help="Do not ask for confirmation before training")
p_train.add_argument("--dry-run", action="store_true", help="List matches but do not train")
args = parser.parse_args(argv)
if args.command == "roots":
roots = discover_roots(args.user)
if not roots:
print("No existing mail roots discovered. Use --root /path/to/Maildir or --root /path/to/mbox.")
return 2
for root in roots:
print(root)
return 0
opts = build_options(args)
roots = resolve_roots(args)
if not roots:
eprint("No mail roots found. Use --root /path/to/Maildir or --user USER.")
return 2
if not any([args.from_, args.to, args.subject, args.any, args.after, args.before]):
eprint("Refusing broad search without filters. Use --from, --to, --subject, --any, --after or --before.")
return 2
matches: list[tuple[Candidate, bytes]] = []
for cand, raw in iter_candidates(roots, opts):
matches.append((cand, raw))
if len(matches) >= args.limit:
break
if not matches:
print("No matching messages found.")
return 1
for idx, (cand, _raw) in enumerate(matches, start=1):
print(format_candidate(cand, idx))
if args.command == "search":
return 0
if args.command == "check":
failures = 0
for idx, (cand, raw) in enumerate(matches, start=1):
rc, out, err = call_rspamc("symbols", raw, cand.path, args.rspamc, args.controller)
print(f"\n--- rspamc symbols for match [{idx}] returncode={rc} ---")
if out.strip():
print(out.strip())
if err.strip():
eprint(err.strip())
if rc != 0:
failures += 1
return 0 if failures == 0 else 3
if args.command == "train":
action = "learn_spam" if args.class_ == "spam" else "learn_ham"
if args.dry_run:
print(f"\nDry-run only. Would call: rspamc {action} <message> for {len(matches)} message(s).")
return 0
if not args.yes:
print(f"\nAbout to train {len(matches)} message(s) as {args.class_.upper()} using rspamc action '{action}'.")
if not confirm("Continue?"):
print("Aborted.")
return 130
failures = 0
successes = 0
for idx, (cand, raw) in enumerate(matches, start=1):
print(f"\nTraining match [{idx}] as {args.class_.upper()}: {cand.subject}")
rc, out, err = call_rspamc(action, raw, cand.path, args.rspamc, args.controller)
if out.strip():
print(out.strip())
if err.strip():
eprint(err.strip())
if rc == 0:
successes += 1
else:
failures += 1
print(f"\nDone. Success: {successes}, failed: {failures}.")
return 0 if failures == 0 else 3
return 2
if __name__ == "__main__":
raise SystemExit(main())