MailRelay/modules/exporter.py

"""Automate proton-mail-export-cli via pexpect.

The Proton Export CLI interactive prompt sequence (observed order):
  1. Email address
  2. Password
  3. Two-factor authentication code   (only if 2FA is enabled)
  4. Mailbox password                  (only if separate mailbox password is set)

After authentication the CLI exports all mail to the export directory.
Each message produces two files:
  {messageID}.eml
  {messageID}.metadata.json

Because the exact prompt strings can vary between CLI versions and account
configurations, each pattern is defined as a regex so minor wording differences
are tolerated.  If the CLI changes its prompt wording, update PROMPTS below.
"""

from pathlib import Path
from typing import Optional

import pexpect

from .logger import get_logger
from .tools import BINARY_PATH as CLI_BINARY_PATH, ensure_export_cli

log = get_logger(__name__)

# ---------------------------------------------------------------------------
# Configurable paths
# ---------------------------------------------------------------------------

EXPORT_DIR = Path(__file__).parent.parent / "data" / "exports"

# ---------------------------------------------------------------------------
# Prompt patterns (case-insensitive regex matched against CLI output)
# ---------------------------------------------------------------------------

PROMPTS = {
    "email":            r"[Ee]mail\s*(address)?[\s:>]+",
    "password":         r"[Pp]assword[\s:>]+",
    "totp":             r"[Tt]wo.factor|[Oo]ne.time|[Tt][Oo][Tt][Pp]|[Aa]uth.*code",
    "mailbox_password": r"[Mm]ailbox\s*[Pp]assword[\s:>]+",
    "done":             r"[Ee]xport\s*(complete|finished|done)|[Ss]uccessfully\s*export",
    "error":            r"[Ee]rror|[Ff]ailed|[Ii]nvalid",
}

# Maximum time to wait for each prompt (seconds)
PROMPT_TIMEOUT = 120
# Total export timeout — large mailboxes can take a while
EXPORT_TIMEOUT = 3600


class ExportError(Exception):
    pass


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def run_export(
    email: str,
    password: str,
    totp_code: str,
    mailbox_password: str = "",
    export_dir: Optional[Path] = None,
) -> Path:
    """Drive proton-mail-export-cli and return the export directory path.

    Raises ExportError on authentication failure or unexpected CLI output.
    """
    out_dir = export_dir or EXPORT_DIR
    out_dir.mkdir(parents=True, exist_ok=True)

    log.info("Starting Proton Mail export...")

    try:
        binary = ensure_export_cli()
    except Exception as exc:
        raise ExportError(str(exc)) from exc

    cmd = [str(binary), "--export-dir", str(out_dir)]
    log.info("Starting Proton export CLI: %s", " ".join(cmd))

    child = pexpect.spawn(
        cmd[0], cmd[1:], timeout=PROMPT_TIMEOUT, encoding="utf-8"
    )
    child.logfile_read = _PexpectLogger(log)

    try:
        _drive_cli(child, email, password, totp_code, mailbox_password)
    except pexpect.TIMEOUT as exc:
        child.close(force=True)
        raise ExportError("Timed out waiting for CLI prompt.") from exc
    except pexpect.EOF as exc:
        output = child.before or ""
        child.close()
        if child.exitstatus and child.exitstatus != 0:
            raise ExportError(
                f"CLI exited with code {child.exitstatus}. Output: {output.strip()}"
            ) from exc
        # EOF after a successful export is normal
        log.info("CLI process finished (EOF).")

    log.info("Export complete. Files in: %s", out_dir)
    return out_dir


def _drive_cli(
    child: pexpect.spawn,
    email: str,
    password: str,
    totp_code: str,
    mailbox_password: str,
) -> None:
    """Respond to each interactive prompt in sequence."""
    patterns = [
        pexpect.TIMEOUT,
        pexpect.EOF,
        PROMPTS["email"],
        PROMPTS["password"],
        PROMPTS["totp"],
        PROMPTS["mailbox_password"],
        PROMPTS["done"],
        PROMPTS["error"],
    ]

    totp_sent = False
    mailbox_sent = False

    while True:
        idx = child.expect(patterns, timeout=PROMPT_TIMEOUT)

        if idx == 0:  # TIMEOUT
            raise pexpect.TIMEOUT("No prompt received within timeout.")

        if idx == 1:  # EOF — process exited
            return

        if idx == 2:  # email prompt
            log.debug("CLI requested email.")
            child.sendline(email)

        elif idx == 3:  # password prompt
            # The CLI may show a password prompt for both the account password
            # and the mailbox password.  We track which we've already sent.
            if not mailbox_sent and mailbox_password and totp_sent:
                log.debug("CLI requested mailbox password.")
                child.sendline(mailbox_password)
                mailbox_sent = True
            else:
                log.debug("CLI requested account password.")
                child.sendline(password)

        elif idx == 4:  # TOTP prompt
            log.debug("CLI requested TOTP code.")
            child.sendline(totp_code)
            totp_sent = True

        elif idx == 5:  # explicit mailbox password prompt
            log.debug("CLI requested mailbox password (explicit prompt).")
            child.sendline(mailbox_password or "")
            mailbox_sent = True

        elif idx == 6:  # export done
            log.info("CLI reported export complete.")
            return

        elif idx == 7:  # error line
            snippet = (child.before or "").strip().splitlines()[-1]
            raise ExportError(f"CLI reported an error: {snippet}")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

class _PexpectLogger:
    """Thin adapter so pexpect writes its read data to our logger at DEBUG."""

    def __init__(self, logger):
        self._log = logger

    def write(self, s: str) -> None:
        if s.strip():
            self._log.debug("[cli] %s", s.rstrip())

    def flush(self) -> None:
        pass