commit 7bc9d7a901e58ff096b482050b6fa0505c43b7e3 Author: Johannes Rest Date: Mon Jun 22 16:17:38 2026 +0200 Initial commit for the rspamd spam/ham trainer diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..85329f2 --- /dev/null +++ b/.editorconfig @@ -0,0 +1,11 @@ +root = true + +[*] +charset = utf-8 +end_of_line = lf +insert_final_newline = true +indent_style = space +indent_size = 4 + +[*.{md,yml,yaml}] +indent_size = 2 diff --git a/.gitea/workflows/ci.yml b/.gitea/workflows/ci.yml new file mode 100644 index 0000000..f575322 --- /dev/null +++ b/.gitea/workflows/ci.yml @@ -0,0 +1,31 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install build tooling + run: | + python -m pip install --upgrade pip + python -m pip install build + + - name: Run tests + run: | + PYTHONPATH=src python -m unittest discover -s tests + + - name: Build package + run: | + python -m build diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ed207c2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,29 @@ +# Python +__pycache__/ +*.py[cod] +*.pyo +*.pyd +.venv/ +venv/ +build/ +dist/ +*.egg-info/ +.pytest_cache/ +.coverage +htmlcov/ + +# Editors +.idea/ +.vscode/ +*.swp +*.swo + +# OS +.DS_Store +Thumbs.db + +# Local mail data, never commit real messages +*.eml +*.mbox +Maildir/ +test-maildata/ diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..45af812 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,8 @@ +# Changelog + +## 0.1.0 + +- Initial Gitea-ready package. +- Maildir and mbox search. +- Spam/Ham training via `rspamc`. +- Defensive filters, confirmation and dry-run mode. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..0fc27e5 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Johannes Rest + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the Software), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED AS IS, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..e187f92 --- /dev/null +++ b/Makefile @@ -0,0 +1,28 @@ +PYTHON ?= python3 + +.PHONY: help test build install uninstall clean + +help: + @echo "Targets:" + @echo " test Run tests" + @echo " build Build source/wheel package" + @echo " install Install package with pip" + @echo " uninstall Remove package" + @echo " clean Remove build artifacts" + +test: + PYTHONPATH=src $(PYTHON) -m unittest discover -s tests + +build: + $(PYTHON) -m pip install --upgrade build + $(PYTHON) -m build + +install: + $(PYTHON) -m pip install . + +uninstall: + $(PYTHON) -m pip uninstall -y rspamd-mail-trainer + +clean: + rm -rf build dist *.egg-info src/*.egg-info .pytest_cache + find . -type d -name __pycache__ -prune -exec rm -rf {} + diff --git a/README.md b/README.md new file mode 100644 index 0000000..92280ca --- /dev/null +++ b/README.md @@ -0,0 +1,151 @@ +# rspamd-mail-trainer + +Small defensive admin tool for Linux mail servers. It finds messages in Maildir +or mbox mailboxes by sender, recipient, subject, date or selected text and then +passes the matching messages to `rspamc learn_spam` or `rspamc learn_ham`. + +Designed for setups such as Rocky Linux + Postfix/Dovecot + Rspamd. + +## Features + +- Searches Maildir and mbox mailboxes. +- Tries to discover common mailbox roots automatically, including `doveconf -n`. +- Filters by `From`, `To`, `Subject`, header/body text and date range. +- Trains spam and ham via `rspamc`. +- Asks for confirmation before training. +- Refuses broad runs without any filter. +- Deduplicates matches by SHA-256. +- Extracts single mbox messages to temporary `.eml` files before training. +- Provides `search`, `check`, `train spam` and `train ham` subcommands. + +## Install + +From the repository root: + +```bash +sudo python3 -m pip install . +``` + +Then: + +```bash +rspamd-mail-trainer --help +``` + +For a local editable install during development: + +```bash +python3 -m pip install -e . +``` + +## Quick start + +Discover mailbox roots for the Unix user `johannes`: + +```bash +sudo rspamd-mail-trainer roots --user johannes +``` + +Search suspicious messages from a sender: + +```bash +sudo rspamd-mail-trainer search --user johannes --from no-reply@example.test +``` + +Train matches as spam: + +```bash +sudo rspamd-mail-trainer train spam --user johannes --from no-reply@example.test +``` + +Train matches as ham: + +```bash +sudo rspamd-mail-trainer train ham --user johannes --subject "Invoice" +``` + +Run Rspamd analysis without training: + +```bash +sudo rspamd-mail-trainer check --user johannes --from no-reply@example.test +``` + +## Useful options + +Treat filter arguments as literal text instead of regex: + +```bash +sudo rspamd-mail-trainer search --user johannes --from no-reply@example.test --literal +``` + +Search subject: + +```bash +sudo rspamd-mail-trainer search --user johannes --subject "Monthly report" +``` + +Search selected headers and text body: + +```bash +sudo rspamd-mail-trainer search --user johannes --any "unsubscribe" +``` + +Restrict by date: + +```bash +sudo rspamd-mail-trainer search --user johannes --from no-reply@example.test --after 2026-06-01 --before 2026-06-13 +``` + +Dry-run training: + +```bash +sudo rspamd-mail-trainer train spam --user johannes --from no-reply@example.test --dry-run +``` + +Skip confirmation for a known-safe batch: + +```bash +sudo rspamd-mail-trainer train spam --user johannes --from no-reply@example.test --yes +``` + +## Push to Gitea + +Create a new empty repository in Gitea, then run from this directory: + +```bash +git init +git add . +git commit -m "Initial import of rspamd mail trainer" +git branch -M main +git remote add origin ssh://git@gitea.example.local/johannes/rspamd-mail-trainer.git +git push -u origin main +``` + +Replace the remote URL with your Gitea repository URL. + +A Gitea Actions workflow is included at `.gitea/workflows/ci.yml`. + +## Development + +Run tests without external dependencies: + +```bash +python3 -m unittest discover -s tests +``` + +Build a Python package: + +```bash +python3 -m pip install build +python3 -m build +``` + +## Safety notes + +This tool reads local mailboxes and can train your Rspamd classifier. Start with +`search`, `check` or `--dry-run`. Avoid broad patterns unless you inspected the +matches. Train good messages as ham from time to time, not only spam. + +## License + +MIT diff --git a/examples/train-ham-subject.sh b/examples/train-ham-subject.sh new file mode 100755 index 0000000..619a811 --- /dev/null +++ b/examples/train-ham-subject.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env sh +set -eu + +sudo rspamd-mail-trainer train ham \ + --user johannes \ + --subject "Invoice" diff --git a/examples/train-spam-sender.sh b/examples/train-spam-sender.sh new file mode 100755 index 0000000..27f0ebc --- /dev/null +++ b/examples/train-spam-sender.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env sh +set -eu + +sudo rspamd-mail-trainer train spam \ + --user johannes \ + --from no-reply@example.test diff --git a/packaging/rpm/rspamd-mail-trainer.spec b/packaging/rpm/rspamd-mail-trainer.spec new file mode 100644 index 0000000..6d12984 --- /dev/null +++ b/packaging/rpm/rspamd-mail-trainer.spec @@ -0,0 +1,39 @@ +Name: rspamd-mail-trainer +Version: 0.1.0 +Release: 1%{?dist} +Summary: Find Maildir/mbox messages and train Rspamd as spam or ham + +License: MIT +URL: https://gitea.example.local/johannes/rspamd-mail-trainer +Source0: %{name}-%{version}.tar.gz + +BuildArch: noarch +BuildRequires: python3-devel +BuildRequires: python3-setuptools +Requires: python3 +Requires: rspamd + +%description +A defensive command line utility for Linux mail servers. It searches Maildir +or mbox mailboxes by sender, recipient, subject, body/header text and date, +then trains Rspamd through rspamc as spam or ham. + +%prep +%autosetup + +%build +%py3_build + +%install +%py3_install + +%files +%license LICENSE +%doc README.md CHANGELOG.md +%{_bindir}/rspamd-mail-trainer +%{python3_sitelib}/rspamd_mail_trainer* +%{python3_sitelib}/rspamd_mail_trainer-%{version}* + +%changelog +* Sat Jun 13 2026 Johannes Rest - 0.1.0-1 +- Initial package diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..d4647a4 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,30 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "rspamd-mail-trainer" +version = "0.1.0" +description = "Find Maildir/mbox messages and train Rspamd as spam or ham." +readme = "README.md" +requires-python = ">=3.9" +license = { text = "MIT" } +authors = [ + { name = "Johannes Rest" } +] +keywords = ["rspamd", "maildir", "mbox", "spam", "ham", "mail"] +classifiers = [ + "Environment :: Console", + "Intended Audience :: System Administrators", + "License :: OSI Approved :: MIT License", + "Operating System :: POSIX :: Linux", + "Programming Language :: Python :: 3", + "Topic :: Communications :: Email :: Filters", + "Topic :: System :: Systems Administration" +] + +[project.scripts] +rspamd-mail-trainer = "rspamd_mail_trainer.cli:main" + +[tool.setuptools.packages.find] +where = ["src"] diff --git a/scripts/install.sh b/scripts/install.sh new file mode 100755 index 0000000..613a38e --- /dev/null +++ b/scripts/install.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env sh +set -eu +python_bin="${PYTHON_BIN:-python3}" +"$python_bin" -m pip install . +echo "Installed rspamd-mail-trainer." +echo "Run: rspamd-mail-trainer --help" diff --git a/scripts/rspamd-mail-trainer b/scripts/rspamd-mail-trainer new file mode 100755 index 0000000..e916fc5 --- /dev/null +++ b/scripts/rspamd-mail-trainer @@ -0,0 +1,4 @@ +#!/usr/bin/env sh +set -eu +PYTHON_BIN="${PYTHON_BIN:-python3}" +exec "$PYTHON_BIN" -m rspamd_mail_trainer "$@" diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh new file mode 100755 index 0000000..03b2042 --- /dev/null +++ b/scripts/uninstall.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env sh +set -eu +python_bin="${PYTHON_BIN:-python3}" +"$python_bin" -m pip uninstall -y rspamd-mail-trainer diff --git a/src/rspamd_mail_trainer/__init__.py b/src/rspamd_mail_trainer/__init__.py new file mode 100644 index 0000000..746fdab --- /dev/null +++ b/src/rspamd_mail_trainer/__init__.py @@ -0,0 +1,3 @@ +"""Rspamd Mail Trainer package.""" + +__version__ = "0.1.0" diff --git a/src/rspamd_mail_trainer/__main__.py b/src/rspamd_mail_trainer/__main__.py new file mode 100644 index 0000000..eb53e2f --- /dev/null +++ b/src/rspamd_mail_trainer/__main__.py @@ -0,0 +1,3 @@ +from .cli import main + +raise SystemExit(main()) diff --git a/src/rspamd_mail_trainer/cli.py b/src/rspamd_mail_trainer/cli.py new file mode 100755 index 0000000..1fd76e5 --- /dev/null +++ b/src/rspamd_mail_trainer/cli.py @@ -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} 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()) diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..705f07d --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,72 @@ +from pathlib import Path +import tempfile +import unittest + +from rspamd_mail_trainer.cli import SearchOptions, iter_candidates, maybe_compile, parse_user_date + + +def write_mail(path: Path, sender: str, subject: str, body: str = "hello") -> None: + path.write_text( + "\n".join([ + f"From: {sender}", + "To: johannes@example.org", + f"Subject: {subject}", + "Date: Fri, 12 Jun 2026 13:36:00 +0000", + "Message-ID: ", + "", + body, + ]), + encoding="utf-8", + ) + + +class CliTests(unittest.TestCase): + def test_maildir_sender_match(self) -> None: + with tempfile.TemporaryDirectory() as td: + root = Path(td) / "Maildir" + cur = root / "cur" + cur.mkdir(parents=True) + write_mail(cur / "1:2,S", "no-reply@example.test", "Monthly report") + opts = SearchOptions( + from_re=maybe_compile("example.test", True, True), + to_re=None, + subject_re=None, + any_re=None, + after=None, + before=None, + max_size=1024 * 1024, + include_seen_duplicates=False, + ) + matches = list(iter_candidates([root], opts)) + self.assertEqual(len(matches), 1) + self.assertEqual(matches[0][0].sender, "no-reply@example.test") + + def test_subject_match(self) -> None: + with tempfile.TemporaryDirectory() as td: + root = Path(td) / "Maildir" + cur = root / "cur" + cur.mkdir(parents=True) + write_mail(cur / "2:2,S", "sender@example.org", "Invoice 123") + opts = SearchOptions( + from_re=None, + to_re=None, + subject_re=maybe_compile("Invoice", True, True), + any_re=None, + after=None, + before=None, + max_size=1024 * 1024, + include_seen_duplicates=False, + ) + matches = list(iter_candidates([root], opts)) + self.assertEqual(len(matches), 1) + + def test_parse_user_date_day_start(self) -> None: + parsed = parse_user_date("2026-06-13") + self.assertIsNotNone(parsed) + assert parsed is not None + self.assertEqual(parsed.hour, 0) + self.assertEqual(parsed.minute, 0) + + +if __name__ == "__main__": + unittest.main()