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

src/rspamd_mail_trainer/cli.py

@@ -29,6 +29,7 @@ import email
 import email.policy
 import getpass
 import hashlib
+import logging
 import mailbox
 import os
 import pwd
@@ -41,10 +42,11 @@ 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
+from typing import Iterator, Optional
 
 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)
@@ -72,14 +74,60 @@ class SearchOptions:
     include_seen_duplicates: bool
 
 
-def eprint(*args: object) -> None:
-    print(*args, file=sys.stderr)
+@dataclass
+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]:
     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:
     if value is None:
         return ""
@@ -268,23 +316,37 @@ def iter_maildir_files(root: Path, max_file_size: int) -> Iterator[Path]:
                 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:
         raw = path.read_bytes()
         msg = parse_message_bytes(raw)
     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
+    if stats:
+        stats.messages_read += 1
     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]]:
+def iter_mbox_candidates(
+    path: Path,
+    opts: SearchOptions,
+    stats: Optional[ScanStats] = None,
+) -> Iterator[tuple[Candidate, bytes]]:
     try:
         mbox = mailbox.mbox(str(path), create=False)
     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
 
     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)
             msg = parse_message_bytes(raw)
         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
+        if stats:
+            stats.messages_read += 1
         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
@@ -303,10 +369,12 @@ def iter_mbox_candidates(path: Path, opts: SearchOptions) -> Iterator[tuple[Cand
 def discover_roots(user: Optional[str]) -> list[Path]:
     roots: list[Path] = []
     user = user or getpass.getuser()
+    LOGGER.info("Discovering mail roots for user %s", user)
 
     # Dovecot mail_location if available.
     doveconf = shutil.which("doveconf")
     if doveconf:
+        LOGGER.debug("Reading Dovecot mail_location via %s -n", doveconf)
         cp = run([doveconf, "-n"])
         if cp.returncode == 0:
             for line in cp.stdout.splitlines():
@@ -362,60 +430,103 @@ def discover_roots(user: Optional[str]) -> list[Path]:
         seen.add(key)
         if rp.exists():
             result.append(rp)
+    LOGGER.info("Discovered %s existing mail root(s)", len(result))
     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}")
+def iter_mbox_files(root: Path, max_file_size: int, seen_paths: set[str]) -> Iterator[Path]:
+    for dirpath, dirnames, filenames in os.walk(root, followlinks=False):
+        p = Path(dirpath)
+        depth = len(p.relative_to(root).parts) if p != root else 0
+        if depth > 3:
+            dirnames[:] = []
             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():
+        dirnames[:] = [
+            d for d in dirnames
+            if d not in {".git", "tmp", "dovecot.index.cache"} and not d.endswith(".lock")
+        ]
+        for filename in filenames:
+            f = p / filename
+            key = path_key(f)
+            if key in seen_paths:
                 continue
             try:
-                if looks_like_mbox(p):
-                    iterator = iter_mbox_candidates(p, opts)
+                st = f.stat()
+            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):
-                    iterator = iter_single_file_candidates(p, opts)
+                    iterator = iter_single_file_candidates(p, opts, stats)
                 else:
+                    stats.non_message_skips += 1
                     continue
 
                 for cand, raw in iterator:
                     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
                     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
             except PermissionError:
-                eprint(f"Permission denied: {p}")
+                stats.unreadable_skips += 1
+                LOGGER.warning("Permission denied: %s", p)
             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:
@@ -453,8 +564,11 @@ def call_rspamc(action: str, raw: bytes, path_hint: Path, rspamc: str, controlle
         if controller:
             cmd.extend(["-h", controller])
         cmd.extend([action, str(target)])
+        LOGGER.info("Running rspamc %s for %s", action, path_hint)
         cp = run(cmd)
         return cp.returncode, cp.stdout, cp.stderr
+    except OSError as exc:
+        return 127, "", f"Cannot run {rspamc}: {exc}"
     finally:
         if temp_path:
             try:
@@ -488,7 +602,13 @@ def resolve_roots(args: argparse.Namespace) -> list[Path]:
     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:
+    add_logging_options(parser)
     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")
@@ -513,6 +633,7 @@ def main(argv: Optional[list[str]] = None) -> int:
     sub = parser.add_subparsers(dest="command", required=True)
 
     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_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")
 
     args = parser.parse_args(argv)
+    configure_logging(args.verbose, args.quiet)
 
     if args.command == "roots":
         roots = discover_roots(args.user)
@@ -546,18 +668,35 @@ def main(argv: Optional[list[str]] = None) -> int:
     roots = resolve_roots(args)
 
     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
 
     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
 
+    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]] = []
-    for cand, raw in iter_candidates(roots, opts):
+    reached_limit = False
+    for cand, raw in iter_candidates(roots, opts, stats):
         matches.append((cand, raw))
         if len(matches) >= args.limit:
+            reached_limit = True
             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:
         print("No matching messages found.")
@@ -577,7 +716,7 @@ def main(argv: Optional[list[str]] = None) -> int:
             if out.strip():
                 print(out.strip())
             if err.strip():
-                eprint(err.strip())
+                LOGGER.warning("%s", err.strip())
             if rc != 0:
                 failures += 1
         return 0 if failures == 0 else 3
@@ -603,7 +742,7 @@ def main(argv: Optional[list[str]] = None) -> int:
             if out.strip():
                 print(out.strip())
             if err.strip():
-                eprint(err.strip())
+                LOGGER.warning("%s", err.strip())
             if rc == 0:
                 successes += 1
             else: