Sanitized mirror from private repository - 2026-04-20 01:32:01 UTC

scripts/proton-organizer/.gitignore

@@ -0,0 +1,2 @@
+config.local.yaml
+processed.db

scripts/proton-organizer/proton_organizer.py

@@ -0,0 +1,330 @@
+#!/usr/bin/env python3
+"""Proton Mail Organizer — classifies emails using a local LLM and applies labels via Proton Bridge."""
+
+import argparse
+import email
+import email.header
+import html
+import imaplib
+import json
+import logging
+import re
+import sqlite3
+import ssl
+import sys
+import urllib.request
+import urllib.error
+from datetime import datetime
+from pathlib import Path
+
+import yaml
+
+LOG_FMT = "%(asctime)s %(levelname)-8s %(message)s"
+log = logging.getLogger("proton-organizer")
+
+DB_PATH = Path(__file__).parent / "processed.db"
+DEFAULT_CONFIG = Path(__file__).parent / "config.local.yaml"
+
+
+# ── helpers ──────────────────────────────────────────────────────────────────
+
+def load_config(path: Path) -> dict:
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+def init_db(db_path: Path) -> sqlite3.Connection:
+    conn = sqlite3.connect(db_path)
+    conn.execute("""
+        CREATE TABLE IF NOT EXISTS processed (
+            message_id TEXT PRIMARY KEY,
+            category TEXT NOT NULL,
+            processed_at TEXT NOT NULL
+        )
+    """)
+    conn.commit()
+    return conn
+
+
+def is_processed(conn: sqlite3.Connection, message_id: str) -> bool:
+    return conn.execute(
+        "SELECT 1 FROM processed WHERE message_id = ?", (message_id,)
+    ).fetchone() is not None
+
+
+def mark_processed(conn: sqlite3.Connection, message_id: str, category: str):
+    conn.execute(
+        "INSERT OR REPLACE INTO processed (message_id, category, processed_at) VALUES (?, ?, ?)",
+        (message_id, category, datetime.now(tz=__import__('zoneinfo').ZoneInfo("UTC")).isoformat()),
+    )
+    conn.commit()
+
+
+def decode_header(raw: str | None) -> str:
+    if not raw:
+        return ""
+    parts = email.header.decode_header(raw)
+    decoded = []
+    for data, charset in parts:
+        if isinstance(data, bytes):
+            try:
+                decoded.append(data.decode(charset or "utf-8", errors="replace"))
+            except (LookupError, UnicodeDecodeError):
+                decoded.append(data.decode("utf-8", errors="replace"))
+        else:
+            decoded.append(data)
+    return " ".join(decoded)
+
+
+def extract_text(msg: email.message.Message, max_chars: int) -> str:
+    body = ""
+    if msg.is_multipart():
+        for part in msg.walk():
+            ct = part.get_content_type()
+            if ct == "text/plain":
+                payload = part.get_payload(decode=True)
+                if payload:
+                    charset = part.get_content_charset() or "utf-8"
+                    body = payload.decode(charset, errors="replace")
+                    break
+            elif ct == "text/html" and not body:
+                payload = part.get_payload(decode=True)
+                if payload:
+                    charset = part.get_content_charset() or "utf-8"
+                    body = html.unescape(re.sub(r"<[^>]+>", " ",
+                                                payload.decode(charset, errors="replace")))
+    else:
+        payload = msg.get_payload(decode=True)
+        if payload:
+            charset = msg.get_content_charset() or "utf-8"
+            body = payload.decode(charset, errors="replace")
+            if msg.get_content_type() == "text/html":
+                body = html.unescape(re.sub(r"<[^>]+>", " ", body))
+
+    return re.sub(r"\s+", " ", body).strip()[:max_chars]
+
+
+# ── Proton Bridge IMAP ──────────────────────────────────────────────────────
+
+class ProtonClient:
+    def __init__(self, email_addr: str, bridge_password: "REDACTED_PASSWORD"
+                 host: str = "127.0.0.1", port: int = 1143):
+        self.email = email_addr
+        ctx = ssl.create_default_context()
+        ctx.check_hostname = False
+        ctx.verify_mode = ssl.CERT_NONE
+        self.conn = imaplib.IMAP4(host, port)
+        self.conn.starttls(ssl_context=ctx)
+        self.conn.login(email_addr, bridge_password)
+
+    def fetch_uids(self, mailbox: str = "INBOX", search: str = "ALL",
+                   batch_size: int = 50) -> list[bytes]:
+        self.conn.select(mailbox)
+        _, data = self.conn.search(None, search)
+        uids = data[0].split()
+        return list(reversed(uids[-batch_size:]))
+
+    def fetch_message(self, uid: bytes) -> email.message.Message:
+        _, data = self.conn.fetch(uid, "(RFC822)")
+        return email.message_from_bytes(data[0][1])
+
+    def apply_label(self, uid: bytes, label: str):
+        """Apply a label by copying the message to the label folder."""
+        try:
+            self.conn.create(label)
+        except imaplib.IMAP4.error:
+            pass
+        result = self.conn.copy(uid, label)
+        if result[0] != "OK":
+            log.warning("Failed to copy to label %s: %s", label, result)
+
+    def archive(self, uid: bytes):
+        """Archive: move from INBOX to Archive folder."""
+        self.conn.copy(uid, "Archive")
+        self.conn.store(uid, "+FLAGS", "\\Deleted")
+        self.conn.expunge()
+
+    def close(self):
+        try:
+            self.conn.close()
+            self.conn.logout()
+        except Exception:
+            pass
+
+
+# ── Ollama LLM ───────────────────────────────────────────────────────────────
+
+def classify_email(ollama_url, model, categories, subject, sender, body_snippet):
+    cat_descriptions = "\n".join(
+        f"- **{name}**: {info['description']}" for name, info in categories.items()
+    )
+    category_names = ", ".join(categories.keys())
+
+    prompt = f"""Classify this email into exactly ONE category. Reply with ONLY the category name, nothing else.
+
+Categories:
+{cat_descriptions}
+
+Email:
+From: {sender}
+Subject: {subject}
+Body: {body_snippet[:1000]}
+
+Reply with one of: {category_names}"""
+
+    payload = json.dumps({
+        "model": model,
+        "prompt": prompt,
+        "stream": False,
+        "options": {"temperature": 0.1, "num_predict": 20},
+    }).encode()
+
+    req = urllib.request.Request(
+        f"{ollama_url.rstrip('/')}/api/generate",
+        data=payload,
+        headers={"Content-Type": "application/json"},
+    )
+
+    try:
+        with urllib.request.urlopen(req, timeout=60) as resp:
+            result = json.loads(resp.read())
+    except urllib.error.URLError as e:
+        log.error("Ollama request failed: %s", e)
+        raise
+
+    raw_response = result.get("response", "").strip().lower()
+    raw_response = re.sub(r"<think>.*?</think>", "", raw_response, flags=re.DOTALL).strip()
+    for name in categories:
+        if name in raw_response:
+            return name
+
+    log.warning("LLM returned unexpected category %r, defaulting to 'personal'", raw_response)
+    return "personal"
+
+
+# ── main ─────────────────────────────────────────────────────────────────────
+
+def run(config_path, dry_run=False, reprocess=False, limit=None):
+    cfg = load_config(config_path)
+    proton_cfg = cfg["proton"]
+    ollama_cfg = cfg["ollama"]
+    categories = cfg["categories"]
+    proc_cfg = cfg.get("processing", {})
+
+    batch_size = limit or proc_cfg.get("batch_size", 50)
+    max_body = proc_cfg.get("max_body_chars", 2000)
+    dry_run = dry_run or proc_cfg.get("dry_run", False)
+    mailbox = proc_cfg.get("mailbox", "INBOX")
+    rules = cfg.get("rules", [])
+
+    log.info("Connecting to Proton Bridge as %s", proton_cfg["email"])
+    client = ProtonClient(
+        proton_cfg["email"],
+        proton_cfg["bridge_password"],
+        host=proton_cfg.get("host", "127.0.0.1"),
+        port=proton_cfg.get("port", 1143),
+    )
+    db = init_db(DB_PATH)
+
+    try:
+        uids = client.fetch_uids(mailbox=mailbox, batch_size=batch_size)
+        log.info("Fetched %d message UIDs", len(uids))
+
+        stats = {cat: 0 for cat in categories}
+        stats["rules"] = 0
+        stats["skipped"] = 0
+        stats["errors"] = 0
+
+        for i, uid in enumerate(uids, 1):
+            try:
+                msg = client.fetch_message(uid)
+                message_id = msg.get("Message-ID", f"uid-{uid.decode()}")
+                subject = decode_header(msg.get("Subject"))
+                sender = decode_header(msg.get("From"))
+
+                if not reprocess and is_processed(db, message_id):
+                    stats["skipped"] += 1
+                    continue
+
+                # Check sender-based rules before LLM
+                rule_matched = False
+                for rule in rules:
+                    for pattern in rule["senders"]:
+                        if pattern.lower() in sender.lower():
+                            folder = rule["folder"]
+                            category = rule.get("category", "personal")
+                            log.info("[%d/%d] Rule match: %s (from: %s) → %s",
+                                     i, len(uids), subject[:60], sender[:40], folder)
+                            if not dry_run:
+                                client.apply_label(uid, folder)
+                                mark_processed(db, message_id, category)
+                            else:
+                                log.info("  [DRY RUN] Would move to: %s", folder)
+                            stats["rules"] += 1
+                            rule_matched = True
+                            break
+                    if rule_matched:
+                        break
+                if rule_matched:
+                    continue
+
+                body = extract_text(msg, max_body)
+                log.info("[%d/%d] Classifying: %s (from: %s)",
+                         i, len(uids), subject[:60], sender[:40])
+
+                category = classify_email(
+                    ollama_cfg["url"], ollama_cfg["model"],
+                    categories, subject, sender, body,
+                )
+                label = categories[category]["label"]
+                log.info("  → %s (%s)", category, label)
+
+                should_archive = categories[category].get("archive", False)
+
+                if not dry_run:
+                    client.apply_label(uid, label)
+                    if should_archive:
+                        client.archive(uid)
+                        log.info("  📥 Archived")
+                    mark_processed(db, message_id, category)
+                else:
+                    log.info("  [DRY RUN] Would apply label: %s%s", label,
+                             " + archive" if should_archive else "")
+
+                stats[category] = stats.get(category, 0) + 1
+
+            except Exception as e:
+                log.error("Error processing UID %s: %s", uid, e)
+                stats["errors"] += 1
+                continue
+
+        log.info("Done! Stats: %s", json.dumps(stats, indent=2))
+
+    finally:
+        client.close()
+        db.close()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Proton Mail Organizer — LLM-powered email classification")
+    parser.add_argument("-c", "--config", type=Path, default=DEFAULT_CONFIG)
+    parser.add_argument("-n", "--dry-run", action="store_true")
+    parser.add_argument("--reprocess", action="store_true")
+    parser.add_argument("--limit", type=int, default=None)
+    parser.add_argument("-v", "--verbose", action="store_true")
+
+    args = parser.parse_args()
+    logging.basicConfig(
+        level=logging.DEBUG if args.verbose else logging.INFO,
+        format=LOG_FMT,
+    )
+
+    if not args.config.exists():
+        log.error("Config not found: %s", args.config)
+        sys.exit(1)
+
+    run(args.config, dry_run=args.dry_run, reprocess=args.reprocess, limit=args.limit)
+
+
+if __name__ == "__main__":
+    main()

scripts/proton-organizer/requirements.txt

@@ -0,0 +1 @@
+pyyaml>=6.0