"""Remote IMAP → local Maildir mailbox import with queued jobs and live progress. Transfers every folder (including Sent) and every message with flags/dates preserved. Jobs are stored in ``mailbox_import_jobs`` and executed by a background worker thread. """ from __future__ import annotations import email.utils import imaplib import logging import os import re import socket import ssl import threading import time from dataclasses import dataclass from datetime import UTC, datetime from email.header import decode_header, make_header from pathlib import Path from typing import Any, Callable from sqlalchemy.orm import Session, joinedload from . import models from .crypto import decrypt_token, encrypt_token from .db import SessionLocal from .settings import get_settings from .utils import mailbox_path_for logger = logging.getLogger(__name__) PROTOCOLS = frozenset({"imap", "imaps", "starttls"}) AUTH_MODES = frozenset({"login", "plain", "cram-md5"}) DEFAULT_PORTS = {"imap": 143, "imaps": 993, "starttls": 143} TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled"}) ACTIVE_STATUSES = frozenset({"pending", "running"}) # Map common provider folder names to Dovecot special-use mailbox names. SPECIAL_FOLDER_ALIASES = { "sent": "Sent", "sent items": "Sent", "sent messages": "Sent", "sent mail": "Sent", "drafts": "Drafts", "draft": "Drafts", "trash": "Trash", "deleted items": "Trash", "deleted": "Trash", "junk": "Junk", "spam": "Junk", "bulk mail": "Junk", "junk e-mail": "Junk", "archive": "Archive", "archives": "Archive", "all mail": "Archive", } _LIST_RE = re.compile( r'\((?P[^)]*)\)\s+(?P"[^"]*"|NIL)\s+(?P"([^"]|\\")*"|\{?\d+\}?.*|.+)', re.IGNORECASE, ) _SAFE_FOLDER_RE = re.compile(r"[^\w.@+=\- ]+", re.UNICODE) _worker_lock = threading.Lock() _worker_started = False _worker_stop = threading.Event() _progress_lock = threading.Lock() def utc_now() -> datetime: return datetime.now(UTC).replace(tzinfo=None) def normalize_protocol(value: str) -> str: protocol = (value or "imaps").strip().lower() if protocol in {"ssl", "tls", "imapssl", "imap-ssl", "imap_ssl"}: protocol = "imaps" if protocol in {"start-tls", "imap-starttls", "imap_starttls"}: protocol = "starttls" if protocol not in PROTOCOLS: raise ValueError(f"Unsupported protocol: {value}. Use imap, imaps or starttls.") return protocol def normalize_auth_mode(value: str) -> str: mode = (value or "login").strip().lower().replace("_", "-") if mode in {"clear", "cleartext", "plaintext", "password"}: mode = "login" if mode not in AUTH_MODES: raise ValueError(f"Unsupported auth mode: {value}. Use login, plain or cram-md5.") return mode def default_port_for(protocol: str) -> int: return DEFAULT_PORTS[normalize_protocol(protocol)] def format_bytes(num: int | float | None) -> str: value = float(num or 0) units = ("B", "KB", "MB", "GB", "TB") for unit in units: if abs(value) < 1024.0 or unit == units[-1]: if unit == "B": return f"{int(value)} {unit}" return f"{value:.2f} {unit}" value /= 1024.0 return f"{value:.2f} TB" def format_speed(bps: int | float | None) -> str: return f"{format_bytes(bps)}/s" def _decode_imap_string(raw: bytes | str) -> str: if isinstance(raw, bytes): try: return raw.decode("utf-8") except UnicodeDecodeError: return raw.decode("latin-1", errors="replace") return str(raw) def _unquote_imap_atom(value: str) -> str: value = value.strip() if value.upper() == "NIL": return "" if len(value) >= 2 and value[0] == '"' and value[-1] == '"': return value[1:-1].replace('\\"', '"').replace("\\\\", "\\") return value def parse_list_response(line: bytes | str) -> tuple[str, str, list[str]] | None: """Return (folder_name, delimiter, flags) from an IMAP LIST/LSUB response line.""" text = _decode_imap_string(line).strip() if text.startswith("* "): text = text[2:] # Drop leading LIST/LSUB token. if text.upper().startswith("LIST "): text = text[5:] elif text.upper().startswith("LSUB "): text = text[5:] match = _LIST_RE.match(text.strip()) if not match: # Fallback: last quoted string or last token is the name. quoted = re.findall(r'"((?:\\.|[^"\\])*)"', text) if quoted: name = quoted[-1].replace('\\"', '"').replace("\\\\", "\\") flags = re.findall(r"\\[A-Za-z0-9-]+", text) return name, "/", flags parts = text.rsplit(None, 1) if len(parts) == 2: return parts[1].strip('"'), "/", [] return None flags_raw = match.group("flags") or "" flags = re.findall(r"\\[A-Za-z0-9-]+", flags_raw) delim = _unquote_imap_atom(match.group("delim") or "/") if not delim or delim.upper() == "NIL": delim = "/" name = _unquote_imap_atom(match.group("name") or "") if not name: return None return name, delim, flags def map_special_folder(name: str) -> str: """Map provider folder names to Dovecot special-use names when appropriate.""" cleaned = name.strip().strip('"') if not cleaned: return cleaned # Use leaf name for alias matching. leaf = cleaned.replace(".", "/").split("/")[-1].strip() alias = SPECIAL_FOLDER_ALIASES.get(leaf.lower()) if alias: # Preserve hierarchy prefix when present. if "/" in cleaned.replace(".", "/"): parent = cleaned.replace(".", "/") parent = "/".join(parent.split("/")[:-1]) return f"{parent}/{alias}" if parent else alias return alias return cleaned def imap_folder_to_maildir_rel(folder: str, delimiter: str = "/") -> str | None: """Return relative maildir++ path under mailbox home, or None for INBOX root.""" name = folder.strip() if not name: return None # Normalize hierarchy separators to / if delimiter and delimiter != "/": name = name.replace(delimiter, "/") name = name.replace(".", "/") # Strip INBOX prefix upper = name.upper() if upper == "INBOX": return None if upper.startswith("INBOX/"): name = name[6:] name = map_special_folder(name) # maildir++ uses '.' as hierarchy separator and leading '.' parts = [p for p in name.replace("/", ".").split(".") if p] safe_parts = [] for part in parts: cleaned = _SAFE_FOLDER_RE.sub("_", part).strip(" .") if cleaned: safe_parts.append(cleaned) if not safe_parts: return None return "." + ".".join(safe_parts) def flags_to_maildir(imap_flags: list[str] | set[str] | tuple[str, ...]) -> str: flag_set = {f.lstrip("\\").lower() for f in imap_flags} chars: list[str] = [] if "seen" in flag_set: chars.append("S") if "answered" in flag_set: chars.append("R") if "flagged" in flag_set: chars.append("F") if "deleted" in flag_set: chars.append("T") if "draft" in flag_set: chars.append("D") return "".join(sorted(chars)) def parse_internaldate(value: str | bytes | None) -> float: if not value: return time.time() text = _decode_imap_string(value).strip().strip('"') try: dt = email.utils.parsedate_to_datetime(text) if dt.tzinfo is None: return dt.replace(tzinfo=UTC).timestamp() return dt.timestamp() except (TypeError, ValueError, IndexError, OverflowError): return time.time() def ensure_maildir(path: Path, uid: int, gid: int) -> None: for sub in ("cur", "new", "tmp"): target = path / sub target.mkdir(parents=True, exist_ok=True) try: os.chown(target, uid, gid) except OSError: pass try: os.chown(path, uid, gid) except OSError: pass # Ensure parent chain is owned when newly created. try: for parent in path.parents: if parent == path.anchor or str(parent) in {"/", "/var", "/var/mail"}: break if parent.exists(): try: os.chown(parent, uid, gid) except OSError: pass except OSError: pass def write_maildir_message( mailbox_home: Path, folder_rel: str | None, raw: bytes, flags: list[str], internal_ts: float, uid: int, gid: int, counter: int, ) -> int: """Write a single RFC822 message into Maildir. Returns bytes written.""" folder_path = mailbox_home if folder_rel is None else mailbox_home / folder_rel ensure_maildir(folder_path, uid, gid) size = len(raw) flag_str = flags_to_maildir(flags) hostname = socket.gethostname().split(".")[0] or "host" hostname = re.sub(r"[^A-Za-z0-9._-]", "_", hostname)[:48] now = time.time() uniq = f"{int(internal_ts)}.M{int((now % 1) * 1_000_000):06d}.P{os.getpid()}.Q{counter}.{hostname},S={size}" if flag_str: filename = f"{uniq}:2,{flag_str}" else: filename = f"{uniq}:2," # Seen messages land in cur/; unseen in new/ dest_dir = folder_path / ("cur" if "S" in flag_str else "new") tmp_path = folder_path / "tmp" / filename final_path = dest_dir / filename # Avoid overwriting on resume collisions. if final_path.exists(): filename = f"{uniq}.R{counter}:2,{flag_str}" tmp_path = folder_path / "tmp" / filename final_path = dest_dir / filename fd = os.open(str(tmp_path), os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) try: written = 0 while written < size: n = os.write(fd, raw[written:]) if n <= 0: raise OSError("short write to maildir tmp") written += n os.fsync(fd) finally: os.close(fd) try: os.chown(tmp_path, uid, gid) except OSError: pass try: os.utime(tmp_path, (internal_ts, internal_ts)) except OSError: pass os.rename(tmp_path, final_path) try: os.chown(final_path, uid, gid) os.utime(final_path, (internal_ts, internal_ts)) except OSError: pass return size class ImportCancelled(Exception): """Raised when an operator cancels a running job.""" class ImportConnectionError(Exception): """Remote IMAP connection/auth failure.""" @dataclass class SourceConfig: host: str port: int protocol: str username: str password: str auth_mode: str = "login" def public_dict(self) -> dict[str, Any]: return { "host": self.host, "port": self.port, "protocol": self.protocol, "username": self.username, "auth_mode": self.auth_mode, } def connect_imap(cfg: SourceConfig, timeout: float = 30.0) -> imaplib.IMAP4: host = cfg.host.strip() if not host: raise ImportConnectionError("Source host is required") port = int(cfg.port or default_port_for(cfg.protocol)) protocol = normalize_protocol(cfg.protocol) auth_mode = normalize_auth_mode(cfg.auth_mode) try: if protocol == "imaps": context = ssl.create_default_context() client: imaplib.IMAP4 = imaplib.IMAP4_SSL(host, port, ssl_context=context, timeout=timeout) else: client = imaplib.IMAP4(host, port, timeout=timeout) if protocol == "starttls": context = ssl.create_default_context() client.starttls(ssl_context=context) except (OSError, imaplib.IMAP4.error, ssl.SSLError) as exc: raise ImportConnectionError(f"Unable to connect to {host}:{port} ({protocol}): {exc}") from exc try: if auth_mode == "cram-md5": client.login_cram_md5(cfg.username, cfg.password) elif auth_mode == "plain": # AUTHENTICATE PLAIN: \0user\0pass try: client.authenticate( "PLAIN", lambda _: f"\0{cfg.username}\0{cfg.password}".encode("utf-8"), ) except imaplib.IMAP4.error: client.login(cfg.username, cfg.password) else: client.login(cfg.username, cfg.password) except imaplib.IMAP4.error as exc: try: client.logout() except Exception: pass raise ImportConnectionError(f"Authentication failed for {cfg.username}: {exc}") from exc return client def test_connection(cfg: SourceConfig) -> dict[str, Any]: client = connect_imap(cfg, timeout=20.0) try: typ, data = client.list() folders: list[str] = [] if typ == "OK" and data: for line in data: if not line: continue parsed = parse_list_response(line) if parsed: folders.append(parsed[0]) # Quick INBOX status for size estimate. messages = 0 try: typ, data = client.status("INBOX", "(MESSAGES)") if typ == "OK" and data and data[0]: m = re.search(r"MESSAGES\s+(\d+)", _decode_imap_string(data[0]), re.I) if m: messages = int(m.group(1)) except Exception: pass return { "ok": True, "folders": sorted(set(folders), key=str.lower), "folder_count": len(set(folders)), "inbox_messages": messages, "source": cfg.public_dict(), } finally: try: client.logout() except Exception: pass def list_folders(client: imaplib.IMAP4) -> list[tuple[str, str, list[str]]]: folders: list[tuple[str, str, list[str]]] = [] seen: set[str] = set() for command in ("list", "lsub"): try: typ, data = getattr(client, command)() except imaplib.IMAP4.error: continue if typ != "OK" or not data: continue for line in data: if not line: continue parsed = parse_list_response(line) if not parsed: continue name, delim, flags = parsed if name in seen: continue # Skip \Noselect folders (namespace markers). flag_l = {f.lstrip("\\").lower() for f in flags} if "noselect" in flag_l or "nonexistent" in flag_l: continue seen.add(name) folders.append((name, delim or "/", flags)) if not any(f[0].upper() == "INBOX" for f in folders): folders.insert(0, ("INBOX", "/", [])) # Prefer INBOX first, then alpha. folders.sort(key=lambda item: (0 if item[0].upper() == "INBOX" else 1, item[0].lower())) return folders def _folder_uid_sizes(client: imaplib.IMAP4, folder: str) -> list[tuple[str, int]]: """Return list of (uid, size) for all messages in folder.""" try: typ, _ = client.select(f'"{folder}"' if " " in folder or "/" in folder else folder, readonly=True) if typ != "OK": # Retry with raw name. typ, _ = client.select(folder, readonly=True) if typ != "OK": return [] except imaplib.IMAP4.error: return [] try: typ, data = client.uid("search", None, "ALL") except imaplib.IMAP4.error: return [] if typ != "OK" or not data or not data[0]: return [] uids = _decode_imap_string(data[0]).split() if not uids: return [] results: list[tuple[str, int]] = [] # Batch FETCH RFC822.SIZE batch_size = 200 for i in range(0, len(uids), batch_size): batch = uids[i : i + batch_size] uid_set = ",".join(batch) try: typ, data = client.uid("fetch", uid_set, "(RFC822.SIZE)") except imaplib.IMAP4.error: for uid in batch: results.append((uid, 0)) continue if typ != "OK" or not data: for uid in batch: results.append((uid, 0)) continue # Map responses: b'1 (UID 12 RFC822.SIZE 3456)' etc. found: dict[str, int] = {} for item in data: if not item or not isinstance(item, (bytes, bytearray, str)): continue text = _decode_imap_string(item) uid_m = re.search(r"UID\s+(\d+)", text, re.I) size_m = re.search(r"RFC822\.SIZE\s+(\d+)", text, re.I) if uid_m: found[uid_m.group(1)] = int(size_m.group(1)) if size_m else 0 for uid in batch: results.append((uid, found.get(uid, 0))) return results def _select_folder(client: imaplib.IMAP4, folder: str) -> bool: candidates = [] if " " in folder or any(c in folder for c in ('"', "\\")): candidates.append(f'"{folder.replace(chr(92), chr(92)+chr(92)).replace(chr(34), chr(92)+chr(34))}"') candidates.append(folder) if not folder.startswith('"'): candidates.append(f'"{folder}"') for name in candidates: try: typ, _ = client.select(name, readonly=True) if typ == "OK": return True except imaplib.IMAP4.error: continue return False def _fetch_message(client: imaplib.IMAP4, uid: str) -> tuple[bytes, list[str], float] | None: try: typ, data = client.uid("fetch", uid, "(RFC822 FLAGS INTERNALDATE)") except imaplib.IMAP4.error as exc: logger.warning("FETCH UID %s failed: %s", uid, exc) return None if typ != "OK" or not data: return None raw: bytes | None = None flags: list[str] = [] internal_ts = time.time() for item in data: if isinstance(item, tuple) and len(item) >= 2: meta = _decode_imap_string(item[0]) body = item[1] if isinstance(body, bytes): raw = body elif isinstance(body, str): raw = body.encode("utf-8", errors="replace") flag_m = re.search(r"FLAGS\s*\(([^)]*)\)", meta, re.I) if flag_m: flags = re.findall(r"\\?[A-Za-z0-9$._-]+", flag_m.group(1)) date_m = re.search(r'INTERNALDATE\s+"([^"]+)"', meta, re.I) if date_m: internal_ts = parse_internaldate(date_m.group(1)) elif isinstance(item, (bytes, bytearray, str)): text = _decode_imap_string(item) flag_m = re.search(r"FLAGS\s*\(([^)]*)\)", text, re.I) if flag_m: flags = re.findall(r"\\?[A-Za-z0-9$._-]+", flag_m.group(1)) date_m = re.search(r'INTERNALDATE\s+"([^"]+)"', text, re.I) if date_m: internal_ts = parse_internaldate(date_m.group(1)) if raw is None: return None return raw, flags, internal_ts def serialize_job(job: models.MailboxImportJob) -> dict[str, Any]: bytes_total = int(job.bytes_total or 0) bytes_done = int(job.bytes_done or 0) remaining = max(0, bytes_total - bytes_done) speed = int(job.speed_bps or 0) eta = job.eta_seconds if eta is None and speed > 0 and remaining > 0: eta = int(remaining / speed) return { "id": job.id, "destination_account_id": job.destination_account_id, "destination_username": job.destination_username, "source_host": job.source_host, "source_port": job.source_port, "source_protocol": job.source_protocol, "source_username": job.source_username, "source_auth_mode": job.source_auth_mode, "status": job.status, "cancel_requested": bool(job.cancel_requested), "progress_percent": float(job.progress_percent or 0), "bytes_total": bytes_total, "bytes_done": bytes_done, "bytes_remaining": remaining, "bytes_total_human": format_bytes(bytes_total), "bytes_done_human": format_bytes(bytes_done), "bytes_remaining_human": format_bytes(remaining), "messages_total": int(job.messages_total or 0), "messages_done": int(job.messages_done or 0), "folders_total": int(job.folders_total or 0), "folders_done": int(job.folders_done or 0), "current_folder": job.current_folder, "speed_bps": speed, "speed_human": format_speed(speed), "eta_seconds": eta, "eta_human": _format_eta(eta), "last_error": job.last_error, "status_message": job.status_message, "created_at": job.created_at.isoformat() + "Z" if job.created_at else None, "started_at": job.started_at.isoformat() + "Z" if job.started_at else None, "finished_at": job.finished_at.isoformat() + "Z" if job.finished_at else None, "updated_at": job.updated_at.isoformat() + "Z" if job.updated_at else None, } def _format_eta(seconds: int | None) -> str | None: if seconds is None: return None seconds = max(0, int(seconds)) if seconds < 60: return f"{seconds}s" minutes, sec = divmod(seconds, 60) if minutes < 60: return f"{minutes}m {sec}s" hours, minutes = divmod(minutes, 60) if hours < 48: return f"{hours}h {minutes}m" days, hours = divmod(hours, 24) return f"{days}d {hours}h" def create_job( db: Session, *, account: models.Account, host: str, port: int | None, protocol: str, username: str, password: str, auth_mode: str = "login", ) -> models.MailboxImportJob: protocol_n = normalize_protocol(protocol) auth_n = normalize_auth_mode(auth_mode) port_n = int(port or default_port_for(protocol_n)) if port_n < 1 or port_n > 65535: raise ValueError("Port must be between 1 and 65535") host_n = (host or "").strip() user_n = (username or "").strip() if not host_n: raise ValueError("Source host is required") if not user_n: raise ValueError("Source username is required") if not password: raise ValueError("Source password is required") if len(host_n) > 255: raise ValueError("Source host is too long") if len(user_n) > 320: raise ValueError("Source username is too long") dest_username = account.username enc = encrypt_token(password) if not enc: raise ValueError("Unable to store source password securely") job = models.MailboxImportJob( destination_account_id=account.id, destination_username=dest_username, source_host=host_n, source_port=port_n, source_protocol=protocol_n, source_username=user_n, source_password_enc=enc, source_auth_mode=auth_n, status="pending", status_message="Queued for import", ) db.add(job) db.commit() db.refresh(job) ensure_worker_started() return job def get_job(db: Session, job_id: int) -> models.MailboxImportJob | None: return db.query(models.MailboxImportJob).filter(models.MailboxImportJob.id == job_id).first() def list_jobs(db: Session, *, limit: int = 100, status: str | None = None) -> list[models.MailboxImportJob]: q = db.query(models.MailboxImportJob).order_by(models.MailboxImportJob.id.desc()) if status: q = q.filter(models.MailboxImportJob.status == status) return q.limit(max(1, min(limit, 500))).all() def request_cancel(db: Session, job: models.MailboxImportJob) -> models.MailboxImportJob: if job.status in TERMINAL_STATUSES: return job if job.status == "pending": job.status = "cancelled" job.cancel_requested = True job.finished_at = utc_now() job.status_message = "Cancelled before start" else: job.cancel_requested = True job.status_message = "Cancel requested…" db.add(job) db.commit() db.refresh(job) return job def delete_job(db: Session, job: models.MailboxImportJob) -> None: if job.status == "running": raise ValueError("Cannot delete a running job; cancel it first") db.delete(job) db.commit() def ensure_schema(db: Session | None = None) -> None: """Best-effort create table if missing (dev / upgrades without reinstall).""" from sqlalchemy import text owns = db is None session = db or SessionLocal() try: session.execute( text( """ CREATE TABLE IF NOT EXISTS mailbox_import_jobs ( id BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY, destination_account_id INT UNSIGNED NOT NULL, destination_username VARCHAR(320) NOT NULL, source_host VARCHAR(255) NOT NULL, source_port INT UNSIGNED NOT NULL DEFAULT 993, source_protocol VARCHAR(16) NOT NULL DEFAULT 'imaps', source_username VARCHAR(320) NOT NULL, source_password_enc TEXT NOT NULL, source_auth_mode VARCHAR(16) NOT NULL DEFAULT 'login', status VARCHAR(16) NOT NULL DEFAULT 'pending', cancel_requested TINYINT(1) NOT NULL DEFAULT 0, progress_percent DOUBLE NOT NULL DEFAULT 0, bytes_total BIGINT UNSIGNED NOT NULL DEFAULT 0, bytes_done BIGINT UNSIGNED NOT NULL DEFAULT 0, messages_total BIGINT UNSIGNED NOT NULL DEFAULT 0, messages_done BIGINT UNSIGNED NOT NULL DEFAULT 0, folders_total INT UNSIGNED NOT NULL DEFAULT 0, folders_done INT UNSIGNED NOT NULL DEFAULT 0, current_folder VARCHAR(512) DEFAULT NULL, speed_bps BIGINT UNSIGNED NOT NULL DEFAULT 0, eta_seconds INT UNSIGNED DEFAULT NULL, last_error VARCHAR(1024) DEFAULT NULL, status_message VARCHAR(512) DEFAULT NULL, resume_state JSON DEFAULT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, started_at TIMESTAMP NULL DEFAULT NULL, finished_at TIMESTAMP NULL DEFAULT NULL, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, KEY idx_import_jobs_status (status, created_at), KEY idx_import_jobs_dest (destination_account_id) ) ENGINE=InnoDB """ ) ) session.commit() except Exception as exc: session.rollback() logger.debug("ensure_schema mailbox_import_jobs: %s", exc) finally: if owns: session.close() def _claim_next_job_id() -> int | None: db = SessionLocal() try: job = ( db.query(models.MailboxImportJob) .filter(models.MailboxImportJob.status == "pending") .order_by(models.MailboxImportJob.id.asc()) .with_for_update(skip_locked=True) .first() ) if not job: # Also reclaim stuck? skip for now. return None if job.cancel_requested: job.status = "cancelled" job.finished_at = utc_now() job.status_message = "Cancelled before start" db.add(job) db.commit() return None job.status = "running" job.started_at = utc_now() job.status_message = "Connecting to source IMAP…" job.progress_percent = 0.0 db.add(job) db.commit() return int(job.id) except Exception as exc: db.rollback() # MySQL without skip_locked / SQLite tests: fallback without lock. logger.debug("claim with lock failed (%s); retrying without lock", exc) try: job = ( db.query(models.MailboxImportJob) .filter(models.MailboxImportJob.status == "pending") .order_by(models.MailboxImportJob.id.asc()) .first() ) if not job: return None if job.cancel_requested: job.status = "cancelled" job.finished_at = utc_now() job.status_message = "Cancelled before start" db.add(job) db.commit() return None job.status = "running" job.started_at = utc_now() job.status_message = "Connecting to source IMAP…" db.add(job) db.commit() return int(job.id) except Exception as exc2: db.rollback() logger.warning("Unable to claim import job: %s", exc2) return None finally: db.close() def _is_cancel_requested(job_id: int) -> bool: db = SessionLocal() try: row = db.query(models.MailboxImportJob.cancel_requested, models.MailboxImportJob.status).filter( models.MailboxImportJob.id == job_id ).first() if not row: return True return bool(row[0]) or row[1] == "cancelled" finally: db.close() def _update_job(job_id: int, **fields: Any) -> None: db = SessionLocal() try: job = db.query(models.MailboxImportJob).filter(models.MailboxImportJob.id == job_id).first() if not job: return for key, value in fields.items(): if hasattr(job, key): setattr(job, key, value) job.updated_at = utc_now() db.add(job) db.commit() except Exception as exc: db.rollback() logger.warning("Failed to update import job %s: %s", job_id, exc) finally: db.close() def run_job(job_id: int) -> None: """Execute a single import job (blocking). Safe to call from worker or CLI.""" settings = get_settings() uid = int(settings.vmail_uid) gid = int(settings.vmail_gid) db = SessionLocal() try: job = ( db.query(models.MailboxImportJob) .options(joinedload(models.MailboxImportJob.destination_account).joinedload(models.Account.domain)) .filter(models.MailboxImportJob.id == job_id) .first() ) if not job: return if job.status == "cancelled": return account = job.destination_account if not account or not account.domain: _update_job(job_id, status="failed", finished_at=utc_now(), last_error="Destination account missing") return password = decrypt_token(job.source_password_enc) or "" cfg = SourceConfig( host=job.source_host, port=int(job.source_port), protocol=job.source_protocol, username=job.source_username, password=password, auth_mode=job.source_auth_mode, ) mailbox_home = mailbox_path_for(account.domain.name, account.local_part) finally: db.close() client: imaplib.IMAP4 | None = None counter = 0 bytes_done = 0 messages_done = 0 folders_done = 0 window_bytes = 0 window_start = time.monotonic() last_persist = 0.0 def persist(force: bool = False, **extra: Any) -> None: nonlocal last_persist, window_bytes, window_start now = time.monotonic() if not force and (now - last_persist) < 0.75: return elapsed = max(0.001, now - window_start) speed = int(window_bytes / elapsed) if window_bytes else 0 # EMA-ish: blend with previous speed in DB via extra if provided if "speed_bps" not in extra: extra["speed_bps"] = speed bytes_total = int(extra.get("bytes_total", 0) or 0) # Prefer explicit totals if set on job via extra fields only for percent b_total = extra.get("_bytes_total_for_pct") m_total = extra.get("_messages_total_for_pct") b_total = int(b_total) if b_total is not None else None m_total = int(m_total) if m_total is not None else None # Compute percent from messages when sizes unknown. pct = 0.0 if b_total and b_total > 0: pct = min(100.0, (bytes_done / b_total) * 100.0) elif m_total and m_total > 0: pct = min(100.0, (messages_done / m_total) * 100.0) extra.pop("_bytes_total_for_pct", None) extra.pop("_messages_total_for_pct", None) remaining = max(0, (b_total or 0) - bytes_done) if b_total else 0 eta = int(remaining / speed) if speed > 0 and remaining > 0 else None # Locals win for counters; callers may still pass status/current_folder/etc. extra.pop("bytes_done", None) extra.pop("messages_done", None) extra.pop("folders_done", None) if "progress_percent" not in extra: extra["progress_percent"] = round(pct, 2) if "eta_seconds" not in extra and eta is not None: extra["eta_seconds"] = eta elif "eta_seconds" not in extra: extra["eta_seconds"] = eta _update_job( job_id, bytes_done=bytes_done, messages_done=messages_done, folders_done=folders_done, **extra, ) last_persist = now # Reset speed window every few seconds so rate stays current. if elapsed >= 3.0: window_bytes = 0 window_start = now try: ensure_maildir(mailbox_home, uid, gid) _update_job(job_id, status_message="Connecting to source IMAP…", current_folder=None) if _is_cancel_requested(job_id): raise ImportCancelled() client = connect_imap(cfg, timeout=45.0) _update_job(job_id, status_message="Listing folders…") folders = list_folders(client) folders_total = len(folders) _update_job(job_id, folders_total=folders_total, status_message=f"Scanning {folders_total} folders…") # Phase 1: inventory sizes inventory: list[tuple[str, str, list[tuple[str, int]]]] = [] messages_total = 0 bytes_total = 0 for folder_name, delim, _flags in folders: if _is_cancel_requested(job_id): raise ImportCancelled() _update_job(job_id, current_folder=folder_name, status_message=f"Scanning {folder_name}…") if not _select_folder(client, folder_name): logger.warning("Job %s: cannot select folder %s", job_id, folder_name) inventory.append((folder_name, delim, [])) continue uid_sizes = _folder_uid_sizes(client, folder_name) inventory.append((folder_name, delim, uid_sizes)) messages_total += len(uid_sizes) bytes_total += sum(size for _, size in uid_sizes) _update_job( job_id, messages_total=messages_total, bytes_total=bytes_total, status_message=f"Importing {messages_total} messages…", ) # Phase 2: copy messages for folder_name, delim, uid_sizes in inventory: if _is_cancel_requested(job_id): raise ImportCancelled() folder_rel = imap_folder_to_maildir_rel(folder_name, delim) if folder_rel: ensure_maildir(mailbox_home / folder_rel, uid, gid) else: ensure_maildir(mailbox_home, uid, gid) persist( force=True, current_folder=folder_name, status_message=f"Importing {folder_name} ({len(uid_sizes)} messages)…", _bytes_total_for_pct=bytes_total, _messages_total_for_pct=messages_total, ) if not uid_sizes: folders_done += 1 continue if not _select_folder(client, folder_name): folders_done += 1 continue for uid_str, expected_size in uid_sizes: if _is_cancel_requested(job_id): raise ImportCancelled() fetched = _fetch_message(client, uid_str) if not fetched: messages_done += 1 persist( _bytes_total_for_pct=bytes_total, _messages_total_for_pct=messages_total, current_folder=folder_name, ) continue raw, flags, internal_ts = fetched counter += 1 try: written = write_maildir_message( mailbox_home, folder_rel, raw, flags, internal_ts, uid, gid, counter, ) except OSError as exc: raise RuntimeError(f"Failed writing message UID {uid_str} in {folder_name}: {exc}") from exc bytes_done += written or expected_size or len(raw) messages_done += 1 window_bytes += written or len(raw) persist( _bytes_total_for_pct=bytes_total, _messages_total_for_pct=messages_total, current_folder=folder_name, status_message=f"Importing {folder_name}… {messages_done}/{messages_total}", ) folders_done += 1 persist( force=True, folders_done=folders_done, _bytes_total_for_pct=bytes_total, _messages_total_for_pct=messages_total, current_folder=folder_name, ) persist( force=True, status="completed", finished_at=utc_now(), progress_percent=100.0, status_message="Import completed", current_folder=None, speed_bps=0, eta_seconds=0, bytes_done=bytes_done, messages_done=messages_done, folders_done=folders_done, ) logger.info( "Mailbox import job %s completed: %s messages, %s bytes → %s", job_id, messages_done, bytes_done, mailbox_home, ) except ImportCancelled: _update_job( job_id, status="cancelled", finished_at=utc_now(), status_message="Cancelled by operator", cancel_requested=True, speed_bps=0, ) logger.info("Mailbox import job %s cancelled", job_id) except ImportConnectionError as exc: _update_job( job_id, status="failed", finished_at=utc_now(), last_error=str(exc)[:1024], status_message="Connection failed", speed_bps=0, ) logger.warning("Mailbox import job %s connection error: %s", job_id, exc) except Exception as exc: _update_job( job_id, status="failed", finished_at=utc_now(), last_error=str(exc)[:1024], status_message="Import failed", speed_bps=0, ) logger.exception("Mailbox import job %s failed", job_id) finally: if client is not None: try: client.logout() except Exception: try: client.shutdown() except Exception: pass def _run_job_privileged(job_id: int) -> None: """Prefer root helper (sudo) so Maildir ownership is vmail; fall back to in-process.""" if os.geteuid() == 0: run_job(job_id) return # Avoid recursive sudo when the helper already re-entered as root via CLI. if os.getenv("LIMRISTEM_MAIL_IMPORT_INPROCESS", "").lower() in {"1", "true", "yes", "on"}: run_job(job_id) return try: from .admin_ops import run_root_script run_root_script("manage-mailbox-import.sh", "run", str(job_id)) return except Exception as exc: logger.warning( "Privileged import helper unavailable (%s); running job %s in-process " "(Maildir ownership may fail if not root/vmail)", exc, job_id, ) run_job(job_id) def _worker_loop() -> None: logger.info("Mailbox import worker started") while not _worker_stop.is_set(): try: job_id = _claim_next_job_id() if job_id is not None: _run_job_privileged(job_id) continue except Exception: logger.exception("Mailbox import worker iteration failed") _worker_stop.wait(2.0) logger.info("Mailbox import worker stopped") def ensure_worker_started() -> None: global _worker_started with _worker_lock: if _worker_started: return # In tests we may disable the background worker. if os.getenv("LIMRISTEM_MAIL_DISABLE_IMPORT_WORKER", "").lower() in {"1", "true", "yes", "on"}: _worker_started = True return thread = threading.Thread(target=_worker_loop, name="limristem-mailbox-import", daemon=True) thread.start() _worker_started = True def stop_worker_for_tests() -> None: """Signal worker stop (tests only).""" global _worker_started _worker_stop.set() with _worker_lock: _worker_started = False _worker_stop.clear()