Updated codebase to include more logging and also accessing large email directories better.

This commit is contained in:
2026-06-25 08:44:23 +02:00
parent 7bc9d7a901
commit 93099f4d76
3 changed files with 296 additions and 53 deletions
+20 -1
View File
@@ -15,6 +15,7 @@ Designed for setups such as Rocky Linux + Postfix/Dovecot + Rspamd.
- Asks for confirmation before training. - Asks for confirmation before training.
- Refuses broad runs without any filter. - Refuses broad runs without any filter.
- Deduplicates matches by SHA-256. - Deduplicates matches by SHA-256.
- Logs scan progress and summaries to stderr while keeping search results on stdout.
- Extracts single mbox messages to temporary `.eml` files before training. - Extracts single mbox messages to temporary `.eml` files before training.
- Provides `search`, `check`, `train spam` and `train ham` subcommands. - Provides `search`, `check`, `train spam` and `train ham` subcommands.
@@ -108,6 +109,22 @@ Skip confirmation for a known-safe batch:
sudo rspamd-mail-trainer train spam --user johannes --from no-reply@example.test --yes sudo rspamd-mail-trainer train spam --user johannes --from no-reply@example.test --yes
``` ```
Reduce console progress output:
```bash
sudo rspamd-mail-trainer search --quiet --user johannes --from no-reply@example.test
```
Show additional diagnostic details such as duplicate suppression:
```bash
sudo rspamd-mail-trainer search --verbose --user johannes --from no-reply@example.test
```
Progress and warnings are written to stderr. Matching messages and `rspamc`
results are written to stdout, so you can still redirect or pipe search output
without mixing it with progress messages.
## Push to Gitea ## Push to Gitea
Create a new empty repository in Gitea, then run from this directory: Create a new empty repository in Gitea, then run from this directory:
@@ -130,7 +147,9 @@ A Gitea Actions workflow is included at `.gitea/workflows/ci.yml`.
Run tests without external dependencies: Run tests without external dependencies:
```bash ```bash
python3 -m unittest discover -s tests make test
# or:
PYTHONPATH=src python3 -m unittest discover -s tests
``` ```
Build a Python package: Build a Python package:
+190 -51
View File
@@ -29,6 +29,7 @@ import email
import email.policy import email.policy
import getpass import getpass
import hashlib import hashlib
import logging
import mailbox import mailbox
import os import os
import pwd import pwd
@@ -41,10 +42,11 @@ from dataclasses import dataclass
from email.header import decode_header, make_header from email.header import decode_header, make_header
from email.message import Message from email.message import Message
from pathlib import Path from pathlib import Path
from typing import Iterable, Iterator, Optional from typing import Iterator, Optional
DEFAULT_MAX_SIZE = 25 * 1024 * 1024 DEFAULT_MAX_SIZE = 25 * 1024 * 1024
MBOX_NAMES = {"mbox", "johannes"} # used only as a weak hint; actual detection is content-based PROGRESS_EVERY_FILES = 1000
LOGGER = logging.getLogger("rspamd_mail_trainer")
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -72,14 +74,60 @@ class SearchOptions:
include_seen_duplicates: bool include_seen_duplicates: bool
def eprint(*args: object) -> None: @dataclass
print(*args, file=sys.stderr) class ScanStats:
roots_seen: int = 0
files_seen: int = 0
messages_read: int = 0
matches_found: int = 0
duplicate_skips: int = 0
unreadable_skips: int = 0
non_message_skips: int = 0
def configure_logging(verbose: int = 0, quiet: bool = False) -> None:
level = logging.WARNING if quiet else logging.INFO
if verbose:
level = logging.DEBUG
if LOGGER.handlers:
handlers = LOGGER.handlers
else:
handler = logging.StreamHandler()
LOGGER.addHandler(handler)
handlers = LOGGER.handlers
formatter = logging.Formatter("%(levelname)s: %(message)s")
for handler in handlers:
handler.setFormatter(formatter)
if isinstance(handler, logging.StreamHandler):
handler.setStream(sys.stderr)
LOGGER.setLevel(level)
LOGGER.propagate = False
def run(cmd: list[str], check: bool = False) -> subprocess.CompletedProcess[str]: 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) return subprocess.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=check)
def path_key(path: Path) -> str:
try:
return str(path.resolve(strict=False))
except Exception:
return str(path)
def log_scan_progress(stats: ScanStats) -> None:
if stats.files_seen and stats.files_seen % PROGRESS_EVERY_FILES == 0:
LOGGER.info(
"Scan progress: checked %s file(s), read %s message(s), found %s match(es)",
stats.files_seen,
stats.messages_read,
stats.matches_found,
)
def decode_header_value(value: object) -> str: def decode_header_value(value: object) -> str:
if value is None: if value is None:
return "" return ""
@@ -268,23 +316,37 @@ def iter_maildir_files(root: Path, max_file_size: int) -> Iterator[Path]:
yield f yield f
def iter_single_file_candidates(path: Path, opts: SearchOptions) -> Iterator[tuple[Candidate, bytes]]: def iter_single_file_candidates(
path: Path,
opts: SearchOptions,
stats: Optional[ScanStats] = None,
) -> Iterator[tuple[Candidate, bytes]]:
try: try:
raw = path.read_bytes() raw = path.read_bytes()
msg = parse_message_bytes(raw) msg = parse_message_bytes(raw)
except Exception as exc: except Exception as exc:
eprint(f"Skip unreadable mail file: {path}: {exc}") if stats:
stats.unreadable_skips += 1
LOGGER.warning("Skip unreadable mail file: %s: %s", path, exc)
return return
if stats:
stats.messages_read += 1
cand = make_candidate(source_type="maildir-file", path=path, raw=raw, msg=msg) cand = make_candidate(source_type="maildir-file", path=path, raw=raw, msg=msg)
if candidate_matches(cand, msg, opts): if candidate_matches(cand, msg, opts):
yield cand, raw yield cand, raw
def iter_mbox_candidates(path: Path, opts: SearchOptions) -> Iterator[tuple[Candidate, bytes]]: def iter_mbox_candidates(
path: Path,
opts: SearchOptions,
stats: Optional[ScanStats] = None,
) -> Iterator[tuple[Candidate, bytes]]:
try: try:
mbox = mailbox.mbox(str(path), create=False) mbox = mailbox.mbox(str(path), create=False)
except Exception as exc: except Exception as exc:
eprint(f"Skip unreadable mbox: {path}: {exc}") if stats:
stats.unreadable_skips += 1
LOGGER.warning("Skip unreadable mbox: %s: %s", path, exc)
return return
for index, _key in enumerate(mbox.keys()): for index, _key in enumerate(mbox.keys()):
@@ -293,8 +355,12 @@ def iter_mbox_candidates(path: Path, opts: SearchOptions) -> Iterator[tuple[Cand
raw = bytes(m) raw = bytes(m)
msg = parse_message_bytes(raw) msg = parse_message_bytes(raw)
except Exception as exc: except Exception as exc:
eprint(f"Skip unreadable mbox message {path}#{index}: {exc}") if stats:
stats.unreadable_skips += 1
LOGGER.warning("Skip unreadable mbox message %s#%s: %s", path, index, exc)
continue continue
if stats:
stats.messages_read += 1
cand = make_candidate(source_type="mbox-message", path=path, raw=raw, msg=msg, message_index=index) cand = make_candidate(source_type="mbox-message", path=path, raw=raw, msg=msg, message_index=index)
if candidate_matches(cand, msg, opts): if candidate_matches(cand, msg, opts):
yield cand, raw yield cand, raw
@@ -303,10 +369,12 @@ def iter_mbox_candidates(path: Path, opts: SearchOptions) -> Iterator[tuple[Cand
def discover_roots(user: Optional[str]) -> list[Path]: def discover_roots(user: Optional[str]) -> list[Path]:
roots: list[Path] = [] roots: list[Path] = []
user = user or getpass.getuser() user = user or getpass.getuser()
LOGGER.info("Discovering mail roots for user %s", user)
# Dovecot mail_location if available. # Dovecot mail_location if available.
doveconf = shutil.which("doveconf") doveconf = shutil.which("doveconf")
if doveconf: if doveconf:
LOGGER.debug("Reading Dovecot mail_location via %s -n", doveconf)
cp = run([doveconf, "-n"]) cp = run([doveconf, "-n"])
if cp.returncode == 0: if cp.returncode == 0:
for line in cp.stdout.splitlines(): for line in cp.stdout.splitlines():
@@ -362,60 +430,103 @@ def discover_roots(user: Optional[str]) -> list[Path]:
seen.add(key) seen.add(key)
if rp.exists(): if rp.exists():
result.append(rp) result.append(rp)
LOGGER.info("Discovered %s existing mail root(s)", len(result))
return result return result
def iter_candidates(roots: list[Path], opts: SearchOptions) -> Iterator[tuple[Candidate, bytes]]: def iter_mbox_files(root: Path, max_file_size: int, seen_paths: set[str]) -> Iterator[Path]:
seen_digests: set[str] = set() for dirpath, dirnames, filenames in os.walk(root, followlinks=False):
p = Path(dirpath)
for root in roots: depth = len(p.relative_to(root).parts) if p != root else 0
if not root.exists(): if depth > 3:
eprint(f"Skip missing root: {root}") dirnames[:] = []
continue continue
dirnames[:] = [
if root.is_file(): d for d in dirnames
paths = [root] if d not in {".git", "tmp", "dovecot.index.cache"} and not d.endswith(".lock")
else: ]
paths = list(iter_maildir_files(root, opts.max_size)) for filename in filenames:
# Also detect mbox files directly below selected root and a few common mail dirs. f = p / filename
for dirpath, dirnames, filenames in os.walk(root, followlinks=False): key = path_key(f)
depth = len(Path(dirpath).relative_to(root).parts) if Path(dirpath) != root else 0 if key in seen_paths:
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 continue
try: try:
if looks_like_mbox(p): st = f.stat()
iterator = iter_mbox_candidates(p, opts) except OSError:
continue
if not st.st_size or st.st_size > max_file_size:
continue
if looks_like_mbox(f):
yield f
def iter_candidates(
roots: list[Path],
opts: SearchOptions,
stats: Optional[ScanStats] = None,
) -> Iterator[tuple[Candidate, bytes]]:
seen_digests: set[str] = set()
stats = stats or ScanStats()
for root in roots:
stats.roots_seen += 1
if not root.exists():
LOGGER.warning("Skip missing root: %s", root)
continue
LOGGER.info("Scanning mail root: %s", root)
seen_paths: set[str] = set()
if root.is_file():
path_iter: Iterator[tuple[Path, str]] = iter([(root, "auto")])
else:
def walk_paths() -> Iterator[tuple[Path, str]]:
for maildir_path in iter_maildir_files(root, opts.max_size):
yield maildir_path, "maildir"
# Also detect mbox files directly below selected root and a few common mail dirs.
for mbox_path in iter_mbox_files(root, opts.max_size * 200, seen_paths):
yield mbox_path, "mbox"
path_iter = walk_paths()
for p, kind in path_iter:
if not p.exists() or not p.is_file():
continue
key = path_key(p)
if key in seen_paths:
continue
seen_paths.add(key)
stats.files_seen += 1
log_scan_progress(stats)
try:
if kind == "mbox" or (kind == "auto" and looks_like_mbox(p)):
iterator = iter_mbox_candidates(p, opts, stats)
elif looks_like_single_message(p): elif looks_like_single_message(p):
iterator = iter_single_file_candidates(p, opts) iterator = iter_single_file_candidates(p, opts, stats)
else: else:
stats.non_message_skips += 1
continue continue
for cand, raw in iterator: for cand, raw in iterator:
if not opts.include_seen_duplicates and cand.digest in seen_digests: if not opts.include_seen_duplicates and cand.digest in seen_digests:
stats.duplicate_skips += 1
LOGGER.debug("Skip duplicate message sha256=%s path=%s", cand.digest[:12], cand.path)
continue continue
seen_digests.add(cand.digest) seen_digests.add(cand.digest)
stats.matches_found += 1
LOGGER.info(
"Match %s: %s from=%r subject=%r",
stats.matches_found,
cand.source_type,
cand.sender,
cand.subject,
)
yield cand, raw yield cand, raw
except PermissionError: except PermissionError:
eprint(f"Permission denied: {p}") stats.unreadable_skips += 1
LOGGER.warning("Permission denied: %s", p)
except OSError as exc: except OSError as exc:
eprint(f"Cannot read {p}: {exc}") stats.unreadable_skips += 1
LOGGER.warning("Cannot read %s: %s", p, exc)
def format_candidate(c: Candidate, idx: int) -> str: def format_candidate(c: Candidate, idx: int) -> str:
@@ -453,8 +564,11 @@ def call_rspamc(action: str, raw: bytes, path_hint: Path, rspamc: str, controlle
if controller: if controller:
cmd.extend(["-h", controller]) cmd.extend(["-h", controller])
cmd.extend([action, str(target)]) cmd.extend([action, str(target)])
LOGGER.info("Running rspamc %s for %s", action, path_hint)
cp = run(cmd) cp = run(cmd)
return cp.returncode, cp.stdout, cp.stderr return cp.returncode, cp.stdout, cp.stderr
except OSError as exc:
return 127, "", f"Cannot run {rspamc}: {exc}"
finally: finally:
if temp_path: if temp_path:
try: try:
@@ -488,7 +602,13 @@ def resolve_roots(args: argparse.Namespace) -> list[Path]:
return roots return roots
def add_logging_options(parser: argparse.ArgumentParser) -> None:
parser.add_argument("-v", "--verbose", action="count", default=0, help="Show debug-level progress details")
parser.add_argument("--quiet", action="store_true", help="Only show warnings and errors on stderr")
def add_common_filters(parser: argparse.ArgumentParser) -> None: def add_common_filters(parser: argparse.ArgumentParser) -> None:
add_logging_options(parser)
parser.add_argument("--user", help="Unix user whose mailbox should be searched, e.g. johannes") 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("--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("--from", dest="from_", help="Regex or literal match against the From header")
@@ -513,6 +633,7 @@ def main(argv: Optional[list[str]] = None) -> int:
sub = parser.add_subparsers(dest="command", required=True) sub = parser.add_subparsers(dest="command", required=True)
p_roots = sub.add_parser("roots", help="Show discovered mail roots") p_roots = sub.add_parser("roots", help="Show discovered mail roots")
add_logging_options(p_roots)
p_roots.add_argument("--user", help="Unix user, e.g. johannes") p_roots.add_argument("--user", help="Unix user, e.g. johannes")
p_search = sub.add_parser("search", help="Search matching messages; does not train") p_search = sub.add_parser("search", help="Search matching messages; does not train")
@@ -532,6 +653,7 @@ def main(argv: Optional[list[str]] = None) -> int:
p_train.add_argument("--dry-run", action="store_true", help="List matches but do not train") p_train.add_argument("--dry-run", action="store_true", help="List matches but do not train")
args = parser.parse_args(argv) args = parser.parse_args(argv)
configure_logging(args.verbose, args.quiet)
if args.command == "roots": if args.command == "roots":
roots = discover_roots(args.user) roots = discover_roots(args.user)
@@ -546,18 +668,35 @@ def main(argv: Optional[list[str]] = None) -> int:
roots = resolve_roots(args) roots = resolve_roots(args)
if not roots: if not roots:
eprint("No mail roots found. Use --root /path/to/Maildir or --user USER.") LOGGER.error("No mail roots found. Use --root /path/to/Maildir or --user USER.")
return 2 return 2
if not any([args.from_, args.to, args.subject, args.any, args.after, args.before]): 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.") LOGGER.error("Refusing broad search without filters. Use --from, --to, --subject, --any, --after or --before.")
return 2 return 2
LOGGER.info("Using %s root(s)", len(roots))
LOGGER.info("Search limit is %s match(es); max single-message size is %s bytes", args.limit, opts.max_size)
stats = ScanStats()
matches: list[tuple[Candidate, bytes]] = [] matches: list[tuple[Candidate, bytes]] = []
for cand, raw in iter_candidates(roots, opts): reached_limit = False
for cand, raw in iter_candidates(roots, opts, stats):
matches.append((cand, raw)) matches.append((cand, raw))
if len(matches) >= args.limit: if len(matches) >= args.limit:
reached_limit = True
break break
LOGGER.info(
"Scan finished: roots=%s files=%s messages=%s matches=%s duplicates=%s unreadable=%s non_messages=%s",
stats.roots_seen,
stats.files_seen,
stats.messages_read,
stats.matches_found,
stats.duplicate_skips,
stats.unreadable_skips,
stats.non_message_skips,
)
if reached_limit:
LOGGER.info("Stopped after reaching --limit=%s", args.limit)
if not matches: if not matches:
print("No matching messages found.") print("No matching messages found.")
@@ -577,7 +716,7 @@ def main(argv: Optional[list[str]] = None) -> int:
if out.strip(): if out.strip():
print(out.strip()) print(out.strip())
if err.strip(): if err.strip():
eprint(err.strip()) LOGGER.warning("%s", err.strip())
if rc != 0: if rc != 0:
failures += 1 failures += 1
return 0 if failures == 0 else 3 return 0 if failures == 0 else 3
@@ -603,7 +742,7 @@ def main(argv: Optional[list[str]] = None) -> int:
if out.strip(): if out.strip():
print(out.strip()) print(out.strip())
if err.strip(): if err.strip():
eprint(err.strip()) LOGGER.warning("%s", err.strip())
if rc == 0: if rc == 0:
successes += 1 successes += 1
else: else:
+86 -1
View File
@@ -1,8 +1,10 @@
from contextlib import redirect_stderr, redirect_stdout
from io import StringIO
from pathlib import Path from pathlib import Path
import tempfile import tempfile
import unittest import unittest
from rspamd_mail_trainer.cli import SearchOptions, iter_candidates, maybe_compile, parse_user_date from rspamd_mail_trainer.cli import ScanStats, SearchOptions, iter_candidates, main, maybe_compile, parse_user_date
def write_mail(path: Path, sender: str, subject: str, body: str = "hello") -> None: def write_mail(path: Path, sender: str, subject: str, body: str = "hello") -> None:
@@ -41,6 +43,89 @@ class CliTests(unittest.TestCase):
self.assertEqual(len(matches), 1) self.assertEqual(len(matches), 1)
self.assertEqual(matches[0][0].sender, "no-reply@example.test") self.assertEqual(matches[0][0].sender, "no-reply@example.test")
def test_iter_candidates_populates_scan_stats(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,
)
stats = ScanStats()
matches = list(iter_candidates([root], opts, stats))
self.assertEqual(len(matches), 1)
self.assertEqual(stats.roots_seen, 1)
self.assertEqual(stats.files_seen, 1)
self.assertEqual(stats.messages_read, 1)
self.assertEqual(stats.matches_found, 1)
def test_main_logs_scan_progress_to_stderr(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")
stdout = StringIO()
stderr = StringIO()
with redirect_stdout(stdout), redirect_stderr(stderr):
rc = main(["search", "--root", str(root), "--from", "example.test", "--literal"])
self.assertEqual(rc, 0)
self.assertIn("[1] maildir-file", stdout.getvalue())
self.assertIn("INFO: Using 1 root(s)", stderr.getvalue())
self.assertIn("INFO: Scanning mail root:", stderr.getvalue())
self.assertIn("INFO: Scan finished:", stderr.getvalue())
def test_quiet_suppresses_info_logging(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")
stdout = StringIO()
stderr = StringIO()
with redirect_stdout(stdout), redirect_stderr(stderr):
rc = main(["search", "--quiet", "--root", str(root), "--from", "example.test", "--literal"])
self.assertEqual(rc, 0)
self.assertIn("[1] maildir-file", stdout.getvalue())
self.assertNotIn("INFO:", stderr.getvalue())
def test_check_reports_missing_rspamc_without_traceback(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")
stdout = StringIO()
stderr = StringIO()
with redirect_stdout(stdout), redirect_stderr(stderr):
rc = main([
"check",
"--root",
str(root),
"--from",
"example.test",
"--literal",
"--rspamc",
"/no/such/rspamc",
])
self.assertEqual(rc, 3)
self.assertIn("returncode=127", stdout.getvalue())
self.assertIn("Cannot run /no/such/rspamc", stderr.getvalue())
def test_subject_match(self) -> None: def test_subject_match(self) -> None:
with tempfile.TemporaryDirectory() as td: with tempfile.TemporaryDirectory() as td:
root = Path(td) / "Maildir" root = Path(td) / "Maildir"