From 9a342af609b27caeb021f782d4c6e3febdd56623 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 9 Mar 2026 09:42:20 +0000 Subject: [PATCH] fix(imap): normalize CRLF line endings to LF in imap_tasks.py Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com> --- app/tasks/imap_tasks.py | 1206 +++++++++++++++++++-------------------- 1 file changed, 603 insertions(+), 603 deletions(-) diff --git a/app/tasks/imap_tasks.py b/app/tasks/imap_tasks.py index 27c42d86..064ccf06 100644 --- a/app/tasks/imap_tasks.py +++ b/app/tasks/imap_tasks.py @@ -1,603 +1,603 @@ -#!/usr/bin/env python3 -import email -import imaplib -import json -import logging -import os -import re -from datetime import datetime, timedelta, timezone - -import redis -from celery import shared_task - -from app.config import settings -from app.tasks.convert_to_pdf import convert_to_pdf # new conversion task -from app.tasks.process_document import process_document # Updated import -from app.utils.allowed_types import ALLOWED_EXTENSIONS, ALLOWED_MIME_TYPES - -# Database session for per-user IMAP accounts (imported lazily to avoid circular imports) -_db_session_factory = None - -# Maximum length to store as last_error to prevent DB bloat -_MAX_ERROR_LENGTH = 500 - - -def _get_db_session(): - """Return a new SQLAlchemy session (lazy import to avoid startup issues).""" - global _db_session_factory # noqa: PLW0603 - if _db_session_factory is None: - from app.database import SessionLocal - - _db_session_factory = SessionLocal - return _db_session_factory() - - -logger = logging.getLogger(__name__) - -# Initialize Redis connection using Celery's Redis settings -redis_client = redis.StrictRedis.from_url(settings.redis_url, decode_responses=True) - - -def _decrypt_imap_password(password: str | None) -> str | None: - """Decrypt an IMAP account password stored in the database. - - Passwords are stored encrypted (Fernet, ``enc:`` prefix) for new records; - legacy plaintext records are returned unchanged so existing accounts - continue to work until they are next updated via the API. - """ - from app.utils.encryption import decrypt_value - - return decrypt_value(password) - - -LOCK_KEY = "imap_lock" # Unique key for locking -LOCK_EXPIRE = 300 # Lock expires in 5 minutes - -# Local cache file for tracking processed emails -CACHE_FILE = os.path.join(settings.workdir, "processed_mails.json") - - -def acquire_lock(): - """Attempt to acquire a Redis-based lock. If acquired, set an expiration.""" - lock_acquired = redis_client.setnx(LOCK_KEY, "locked") - if lock_acquired: - redis_client.expire(LOCK_KEY, LOCK_EXPIRE) - logger.info("Lock acquired for IMAP processing.") - return True - logger.warning("Lock already held. Skipping this cycle.") - return False - - -def release_lock(): - """Release the lock by deleting the Redis key.""" - redis_client.delete(LOCK_KEY) - logger.info("Lock released.") - - -def load_processed_emails(): - """Load the list of already processed emails from a local JSON file.""" - if os.path.exists(CACHE_FILE): - try: - with open(CACHE_FILE, "r") as f: - processed_emails = json.load(f) - processed_emails = cleanup_old_entries(processed_emails) - return processed_emails - except json.JSONDecodeError: - logger.warning("Failed to decode JSON, resetting processed emails cache.") - return {} - return {} - - -def save_processed_emails(processed_emails): - """Save the processed email IDs to a local JSON file.""" - with open(CACHE_FILE, "w") as f: - json.dump(processed_emails, f, indent=4) - - -def cleanup_old_entries(processed_emails): - """Remove entries older than 7 days from the cache to avoid infinite growth.""" - seven_days_ago = datetime.now(timezone.utc) - timedelta(days=7) - valid_emails = {} - for msg_id, date_str in processed_emails.items(): - naive_dt = datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S") - aware_dt = naive_dt.replace(tzinfo=timezone.utc) - if aware_dt > seven_days_ago: - valid_emails[msg_id] = date_str - return valid_emails - - -@shared_task -def pull_all_inboxes(): - """ - Periodic Celery task that checks all configured IMAP mailboxes - and fetches attachments from new emails. - Ensures only one instance runs at a time using Redis-based locking. - - Processes: - 1. System-level mailboxes configured via environment variables (IMAP1, IMAP2). - 2. Per-user IMAP accounts stored in the ``user_imap_accounts`` database table. - """ - if not acquire_lock(): - logger.info("Skipping execution: Another instance is running.") - return - - try: - logger.info("Starting pull_all_inboxes") - - # Mailbox #1 (non-Gmail) - check_and_pull_mailbox( - mailbox_key="imap1", - host=settings.imap1_host, - port=settings.imap1_port, - username=settings.imap1_username, - password=settings.imap1_password, - use_ssl=settings.imap1_ssl, - delete_after_process=settings.imap1_delete_after_process, - ) - - # Mailbox #2 (Gmail) - check_and_pull_mailbox( - mailbox_key="imap2", - host=settings.imap2_host, - port=settings.imap2_port, - username=settings.imap2_username, - password=settings.imap2_password, - use_ssl=settings.imap2_ssl, - delete_after_process=settings.imap2_delete_after_process, - ) - - # Per-user IMAP accounts from the database - _pull_user_imap_accounts() - - # Per-user IMAP integrations from the UserIntegration model - _pull_user_integration_imap() - - logger.info("Finished pull_all_inboxes") - - finally: - release_lock() - - -def _pull_user_imap_accounts() -> None: - """Iterate over all active per-user IMAP accounts and pull their inboxes.""" - try: - from app.models import UserImapAccount - - db = _get_db_session() - try: - accounts = db.query(UserImapAccount).filter(UserImapAccount.is_active.is_(True)).all() - logger.info("Processing %d per-user IMAP account(s)", len(accounts)) - for acct in accounts: - # Use a descriptive identifier for logging and processed-email cache keys - account_identifier = f"user_{acct.owner_id}_{acct.id}" - try: - pull_inbox( - mailbox_key=account_identifier, - host=acct.host, - port=acct.port, - username=acct.username, - password=_decrypt_imap_password(acct.password), - use_ssl=acct.use_ssl, - delete_after_process=acct.delete_after_process, - owner_id=acct.owner_id, - ) - # Record successful poll - acct.last_checked_at = datetime.now(timezone.utc) - acct.last_error = None - db.commit() - except Exception as exc: # noqa: BLE001 - error_msg = str(exc)[:_MAX_ERROR_LENGTH] - logger.error( - "Error pulling user IMAP account %d (%s@%s): %s", - acct.id, - acct.username, - acct.host, - error_msg, - ) - try: - acct.last_checked_at = datetime.now(timezone.utc) - acct.last_error = error_msg - db.commit() - except Exception: # noqa: BLE001 - db.rollback() - finally: - db.close() - except Exception as exc: # noqa: BLE001 - logger.error("Failed to process per-user IMAP accounts: %s", exc) - - -def _pull_user_integration_imap() -> None: - """Iterate over all active IMAP UserIntegrations and pull their inboxes. - - This polls the ``user_integrations`` table for records with - ``integration_type='IMAP'``, ``direction='SOURCE'``, and ``is_active=True``. - Each integration's config/credentials are decoded and passed to - :func:`pull_inbox` with the owning user's ``owner_id`` so that ingested - documents are correctly attributed. - - Individual connection failures are caught and recorded on the integration - without crashing the polling loop. - """ - try: - import json as _json - - from app.models import IntegrationDirection, IntegrationType, UserIntegration - from app.utils.encryption import decrypt_value - - db = _get_db_session() - try: - integrations = ( - db.query(UserIntegration) - .filter( - UserIntegration.integration_type == IntegrationType.IMAP, - UserIntegration.direction == IntegrationDirection.SOURCE, - UserIntegration.is_active.is_(True), - ) - .all() - ) - logger.info("Processing %d IMAP UserIntegration(s)", len(integrations)) - for integ in integrations: - account_identifier = f"integration_{integ.owner_id}_{integ.id}" - try: - cfg = _json.loads(integ.config) if integ.config else {} - raw_creds = decrypt_value(integ.credentials) if integ.credentials else None - creds = _json.loads(raw_creds) if raw_creds else {} - - host = cfg.get("host") - port = int(cfg.get("port", 993)) - username = cfg.get("username") - password = creds.get("password") - use_ssl = cfg.get("use_ssl", True) - delete_after = cfg.get("delete_after_process", False) - gmail_labels = cfg.get("gmail_apply_labels", True) - - if not (host and username and password): - logger.warning( - "IMAP integration %d (owner %s) has incomplete config — skipping.", - integ.id, - integ.owner_id, - ) - continue - - pull_inbox( - mailbox_key=account_identifier, - host=host, - port=port, - username=username, - password=password, - use_ssl=use_ssl, - delete_after_process=delete_after, - owner_id=integ.owner_id, - gmail_apply_labels=gmail_labels, - ) - integ.last_used_at = datetime.now(timezone.utc) - integ.last_error = None - db.commit() - except Exception as exc: # noqa: BLE001 - error_msg = str(exc)[:_MAX_ERROR_LENGTH] - logger.error( - "Error pulling IMAP integration %d (owner %s): %s", - integ.id, - integ.owner_id, - error_msg, - ) - try: - integ.last_used_at = datetime.now(timezone.utc) - integ.last_error = error_msg - db.commit() - except Exception: # noqa: BLE001 - db.rollback() - finally: - db.close() - except Exception as exc: # noqa: BLE001 - logger.error("Failed to process IMAP UserIntegrations: %s", exc) - - -def check_and_pull_mailbox( - mailbox_key: str, - host: str | None, - port: int | None, - username: str | None, - password: str | None, - use_ssl: bool, - delete_after_process: bool, -): - """Validates config and invokes pulling from the mailbox if valid.""" - if not (host and port and username and password): - logger.warning(f"Mailbox {mailbox_key} is missing config, skipping.") - return - - logger.info(f"Checking mailbox: {mailbox_key}") - pull_inbox( - mailbox_key=mailbox_key, - host=host, - port=port, - username=username, - password=password, - use_ssl=use_ssl, - delete_after_process=delete_after_process, - ) - - -def pull_inbox( - mailbox_key, - host, - port, - username, - password, - use_ssl, - delete_after_process, - owner_id=None, - gmail_apply_labels=True, -): - """ - Connects to the IMAP inbox, fetches new unread emails from the last 3 days, - and processes attachments while preserving the original unread status. - - For Gmail: - - Attempts to select the localized All Mail folder. - - Runs an X-GM-RAW query: "in:anywhere in:unread newer_than:3d has:attachment". - - For non-Gmail mailboxes, it falls back to selecting the INBOX with a SINCE/UNSEEN filter. - - Args: - owner_id: Optional user identifier. When provided, ingested documents are - attributed to this user via ``process_document`` / ``convert_to_pdf``. - gmail_apply_labels: Whether to apply Gmail-specific labels and stars to - processed emails. Only relevant for Gmail hosts. Defaults to True. - """ - logger.info("Connecting to %s at %s:%s (SSL=%s)", mailbox_key, host, port, use_ssl) - processed_emails = load_processed_emails() - - try: - mail = imaplib.IMAP4_SSL(host, port) if use_ssl else imaplib.IMAP4(host, port) - mail.login(username, password) - - is_gmail_host = "gmail" in host.lower() - if is_gmail_host: - # For Gmail, try to select the localized All Mail folder. - all_mail_folder = find_all_mail_folder(mail) - if all_mail_folder: - logger.info("Using Gmail All Mail folder: %s", all_mail_folder) - mail.select(f'"{all_mail_folder}"') - else: - logger.warning("Gmail All Mail folder not found, falling back to INBOX.") - mail.select("INBOX") - # Use the X-GM-RAW query for Gmail. - raw_query = "in:anywhere in:unread newer_than:3d has:attachment" - status, search_data = mail.search(None, "X-GM-RAW", f'"{raw_query}"') - else: - # For non-Gmail, select INBOX and use SINCE/UNSEEN query. - mail.select("INBOX") - since_date = (datetime.now(timezone.utc) - timedelta(days=3)).strftime("%d-%b-%Y") - status, search_data = mail.search(None, f"(SINCE {since_date} UNSEEN)") - - if status != "OK": - logger.warning("Search failed on mailbox %s. Status=%s", mailbox_key, status) - mail.close() - mail.logout() - return - - msg_numbers = search_data[0].split() - logger.info("Found %d unread emails in %s.", len(msg_numbers), mailbox_key) - - for num in msg_numbers: - status, msg_data = mail.fetch(num, "(RFC822)") - if status != "OK": - logger.warning("Failed to fetch message %s in %s. Status=%s", num, mailbox_key, status) - continue - - raw_email = msg_data[0][1] - email_message = email.message_from_bytes(raw_email) - msg_id = email_message.get("Message-ID") - - if not msg_id: - logger.warning("Skipping email without Message-ID in %s", mailbox_key) - continue - - if msg_id in processed_emails: - logger.info("Skipping already processed email %s in %s", msg_id, mailbox_key) - continue - - # For Gmail, check if the email already has the "Ingested" label. - if is_gmail_host and gmail_apply_labels: - if email_already_has_label(mail, num, "Ingested"): - logger.info("Skipping email %s in %s, already labeled 'Ingested'.", msg_id, mailbox_key) - continue - - # Process attachments (and convert non-PDF files). - # We call the function without assigning its return value since it is not used. - fetch_attachments_and_enqueue(email_message, owner_id=owner_id) - - if settings.imap_readonly_mode: - logger.info("Readonly mode: skipping mailbox modifications for %s in %s", msg_id, mailbox_key) - else: - if is_gmail_host and gmail_apply_labels: - mark_as_processed_with_star(mail, num) - mark_as_processed_with_label(mail, num, label="Ingested") - - if delete_after_process: - logger.info("Deleting message %s from %s", num.decode(), mailbox_key) - mail.store(num, "+FLAGS", "\\Deleted") - else: - mail.store(num, "-FLAGS", "\\Seen") - - processed_emails[msg_id] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S") - save_processed_emails(processed_emails) - - if not settings.imap_readonly_mode and delete_after_process: - mail.expunge() - - mail.close() - mail.logout() - logger.info("Finished processing mailbox %s", mailbox_key) - - except Exception as e: - logger.exception("Error pulling mailbox %s: %s", mailbox_key, e) - - -def fetch_attachments_and_enqueue(email_message, owner_id: str | None = None): - """ - Extracts attachments from the email and processes only allowed file types. - - Files are accepted if either: - 1. They have a MIME type from the ALLOWED_MIME_TYPES set, OR - 2. They have a '.pdf' file extension (regardless of MIME type) - - Allowed file types include: - - PDF: application/pdf or *.pdf extension - - Microsoft Office files: - - Word: application/msword, - application/vnd.openxmlformats-officedocument.wordprocessingml.document - - Excel: application/vnd.ms-excel, - application/vnd.openxmlformats-officedocument.spreadsheetml.sheet - - PowerPoint: application/vnd.ms-powerpoint, - application/vnd.openxmlformats-officedocument.presentationml.presentation - - Other meaningful attachments: - - Plain text: text/plain - - CSV: text/csv - - Rich Text Format: application/rtf, text/rtf - - If the attachment is a PDF (by extension or MIME type), it is enqueued for upload; - any other allowed file is enqueued for conversion to PDF. - - Args: - email_message: The parsed email message to extract attachments from. - owner_id: Optional user identifier forwarded to ``process_document`` / - ``convert_to_pdf`` for multi-tenant attribution. - - Returns True if at least one allowed attachment was processed. - """ - has_attachment = False - for part in email_message.walk(): - if part.get_content_maintype() == "multipart": - continue - - filename = part.get_filename() - if not filename: - continue - - # Check if it's a PDF file by extension, regardless of MIME type - is_pdf_by_extension = filename.lower().endswith(".pdf") - - mime_type = part.get_content_type() - file_ext = os.path.splitext(filename)[1].lower() - # Accept file if it has an allowed MIME type, an allowed extension, OR is a PDF by extension - if mime_type not in ALLOWED_MIME_TYPES and file_ext not in ALLOWED_EXTENSIONS and not is_pdf_by_extension: - logger.info("Skipping attachment %s with MIME type %s", filename, mime_type) - continue - - file_path = os.path.join(settings.workdir, filename) - with open(file_path, "wb") as f: - f.write(part.get_payload(decode=True)) - - # If it's a PDF by MIME type or extension, process it directly - if mime_type == "application/pdf" or is_pdf_by_extension: - process_document.delay(file_path, owner_id=owner_id) - logger.info("Enqueued PDF for upload: %s (MIME: %s)", filename, mime_type) - elif mime_type in ALLOWED_MIME_TYPES: - # Other allowed files are sent for conversion - convert_to_pdf.delay(file_path, owner_id=owner_id) - logger.info("Enqueued file for conversion to PDF: %s", filename) - - has_attachment = True - return has_attachment - - -def email_already_has_label(mail, msg_id, label="Ingested"): - """ - Checks if the given message (msg_id) has the specified Gmail label. - Returns True if the label is found, False otherwise. - """ - try: - # Convert msg_id to bytes if it's an integer - if isinstance(msg_id, int): - msg_id = str(msg_id).encode() - - label_status, label_data = mail.fetch(msg_id, "(X-GM-LABELS)") - if label_status == "OK" and label_data and len(label_data) > 0: - raw_labels = label_data[0][1].decode("utf-8", errors="ignore") - if label in raw_labels: - return True - except Exception as e: - logger.error("Failed to fetch labels for msg_id=%s: %s", msg_id, e) - return False - - -def mark_as_processed_with_star(mail, msg_id): - """Stars the email in Gmail.""" - try: - mail.store(msg_id, "+FLAGS", "\\Flagged") - logger.info("Email %s starred in Gmail.", msg_id) - except Exception as e: - logger.error("Failed to star email %s: %s", msg_id, e) - - -def mark_as_processed_with_label(mail, msg_id, label="Ingested"): - """Adds a custom label to the email in Gmail.""" - try: - mail.store(msg_id, "+X-GM-LABELS", label) - logger.info("Email %s labeled '%s' in Gmail.", msg_id, label) - except Exception as e: - logger.error("Failed to label email %s with %s: %s", msg_id, label, e) - - -def find_all_mail_folder(mail): - """ - Attempts to select the Gmail All Mail folder using known localized names. - Falls back to using XLIST if needed. - Returns the folder name if found, otherwise None. - """ - COMMON_ALL_MAIL_NAMES = [ - "[Gmail]/Alle Nachrichten", - "[Gmail]/All Mail", - "[Gmail]/Todos", - "[Gmail]/Tutte le mail", - "[Gmail]/Tous les messages", - ] - for candidate in COMMON_ALL_MAIL_NAMES: - status, _ = mail.select(f'"{candidate}"', readonly=True) - if status == "OK": - return candidate - - capabilities = get_capabilities(mail) - if "XLIST" in capabilities: - candidate = find_all_mail_xlist(mail) - if candidate: - return candidate - return None - - -def get_capabilities(mail): - """Returns a list of capabilities supported by the IMAP server.""" - typ, data = mail.capability() - if typ == "OK" and data: - caps = data[0].decode("utf-8", errors="ignore").upper().split() - return caps - return [] - - -def find_all_mail_xlist(mail): - """ - Uses XLIST to discover the mailbox flagged as All Mail. - Returns the folder name if found, otherwise None. - """ - tag = mail._new_tag().decode("ascii") - command_str = f'{tag} XLIST "" "*"' - mail.send((command_str + "\r\n").encode("utf-8")) - - all_mail_folder = None - while True: - line = mail.readline() - if not line: - break - line_str = line.decode("utf-8", errors="ignore").strip() - if line_str.upper().startswith("* XLIST ") and "\\ALLMAIL" in line_str.upper(): - match = re.search(r'"([^"]+)"$', line_str) - if match: - candidate = match.group(1) - logger.info("Found All Mail folder via XLIST: %s", candidate) - all_mail_folder = candidate - if line_str.startswith(tag): - break - return all_mail_folder +#!/usr/bin/env python3 +import email +import imaplib +import json +import logging +import os +import re +from datetime import datetime, timedelta, timezone + +import redis +from celery import shared_task + +from app.config import settings +from app.tasks.convert_to_pdf import convert_to_pdf # new conversion task +from app.tasks.process_document import process_document # Updated import +from app.utils.allowed_types import ALLOWED_EXTENSIONS, ALLOWED_MIME_TYPES + +# Database session for per-user IMAP accounts (imported lazily to avoid circular imports) +_db_session_factory = None + +# Maximum length to store as last_error to prevent DB bloat +_MAX_ERROR_LENGTH = 500 + + +def _get_db_session(): + """Return a new SQLAlchemy session (lazy import to avoid startup issues).""" + global _db_session_factory # noqa: PLW0603 + if _db_session_factory is None: + from app.database import SessionLocal + + _db_session_factory = SessionLocal + return _db_session_factory() + + +logger = logging.getLogger(__name__) + +# Initialize Redis connection using Celery's Redis settings +redis_client = redis.StrictRedis.from_url(settings.redis_url, decode_responses=True) + + +def _decrypt_imap_password(password: str | None) -> str | None: + """Decrypt an IMAP account password stored in the database. + + Passwords are stored encrypted (Fernet, ``enc:`` prefix) for new records; + legacy plaintext records are returned unchanged so existing accounts + continue to work until they are next updated via the API. + """ + from app.utils.encryption import decrypt_value + + return decrypt_value(password) + + +LOCK_KEY = "imap_lock" # Unique key for locking +LOCK_EXPIRE = 300 # Lock expires in 5 minutes + +# Local cache file for tracking processed emails +CACHE_FILE = os.path.join(settings.workdir, "processed_mails.json") + + +def acquire_lock(): + """Attempt to acquire a Redis-based lock. If acquired, set an expiration.""" + lock_acquired = redis_client.setnx(LOCK_KEY, "locked") + if lock_acquired: + redis_client.expire(LOCK_KEY, LOCK_EXPIRE) + logger.info("Lock acquired for IMAP processing.") + return True + logger.warning("Lock already held. Skipping this cycle.") + return False + + +def release_lock(): + """Release the lock by deleting the Redis key.""" + redis_client.delete(LOCK_KEY) + logger.info("Lock released.") + + +def load_processed_emails(): + """Load the list of already processed emails from a local JSON file.""" + if os.path.exists(CACHE_FILE): + try: + with open(CACHE_FILE, "r") as f: + processed_emails = json.load(f) + processed_emails = cleanup_old_entries(processed_emails) + return processed_emails + except json.JSONDecodeError: + logger.warning("Failed to decode JSON, resetting processed emails cache.") + return {} + return {} + + +def save_processed_emails(processed_emails): + """Save the processed email IDs to a local JSON file.""" + with open(CACHE_FILE, "w") as f: + json.dump(processed_emails, f, indent=4) + + +def cleanup_old_entries(processed_emails): + """Remove entries older than 7 days from the cache to avoid infinite growth.""" + seven_days_ago = datetime.now(timezone.utc) - timedelta(days=7) + valid_emails = {} + for msg_id, date_str in processed_emails.items(): + naive_dt = datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S") + aware_dt = naive_dt.replace(tzinfo=timezone.utc) + if aware_dt > seven_days_ago: + valid_emails[msg_id] = date_str + return valid_emails + + +@shared_task +def pull_all_inboxes(): + """ + Periodic Celery task that checks all configured IMAP mailboxes + and fetches attachments from new emails. + Ensures only one instance runs at a time using Redis-based locking. + + Processes: + 1. System-level mailboxes configured via environment variables (IMAP1, IMAP2). + 2. Per-user IMAP accounts stored in the ``user_imap_accounts`` database table. + """ + if not acquire_lock(): + logger.info("Skipping execution: Another instance is running.") + return + + try: + logger.info("Starting pull_all_inboxes") + + # Mailbox #1 (non-Gmail) + check_and_pull_mailbox( + mailbox_key="imap1", + host=settings.imap1_host, + port=settings.imap1_port, + username=settings.imap1_username, + password=settings.imap1_password, + use_ssl=settings.imap1_ssl, + delete_after_process=settings.imap1_delete_after_process, + ) + + # Mailbox #2 (Gmail) + check_and_pull_mailbox( + mailbox_key="imap2", + host=settings.imap2_host, + port=settings.imap2_port, + username=settings.imap2_username, + password=settings.imap2_password, + use_ssl=settings.imap2_ssl, + delete_after_process=settings.imap2_delete_after_process, + ) + + # Per-user IMAP accounts from the database + _pull_user_imap_accounts() + + # Per-user IMAP integrations from the UserIntegration model + _pull_user_integration_imap() + + logger.info("Finished pull_all_inboxes") + + finally: + release_lock() + + +def _pull_user_imap_accounts() -> None: + """Iterate over all active per-user IMAP accounts and pull their inboxes.""" + try: + from app.models import UserImapAccount + + db = _get_db_session() + try: + accounts = db.query(UserImapAccount).filter(UserImapAccount.is_active.is_(True)).all() + logger.info("Processing %d per-user IMAP account(s)", len(accounts)) + for acct in accounts: + # Use a descriptive identifier for logging and processed-email cache keys + account_identifier = f"user_{acct.owner_id}_{acct.id}" + try: + pull_inbox( + mailbox_key=account_identifier, + host=acct.host, + port=acct.port, + username=acct.username, + password=_decrypt_imap_password(acct.password), + use_ssl=acct.use_ssl, + delete_after_process=acct.delete_after_process, + owner_id=acct.owner_id, + ) + # Record successful poll + acct.last_checked_at = datetime.now(timezone.utc) + acct.last_error = None + db.commit() + except Exception as exc: # noqa: BLE001 + error_msg = str(exc)[:_MAX_ERROR_LENGTH] + logger.error( + "Error pulling user IMAP account %d (%s@%s): %s", + acct.id, + acct.username, + acct.host, + error_msg, + ) + try: + acct.last_checked_at = datetime.now(timezone.utc) + acct.last_error = error_msg + db.commit() + except Exception: # noqa: BLE001 + db.rollback() + finally: + db.close() + except Exception as exc: # noqa: BLE001 + logger.error("Failed to process per-user IMAP accounts: %s", exc) + + +def _pull_user_integration_imap() -> None: + """Iterate over all active IMAP UserIntegrations and pull their inboxes. + + This polls the ``user_integrations`` table for records with + ``integration_type='IMAP'``, ``direction='SOURCE'``, and ``is_active=True``. + Each integration's config/credentials are decoded and passed to + :func:`pull_inbox` with the owning user's ``owner_id`` so that ingested + documents are correctly attributed. + + Individual connection failures are caught and recorded on the integration + without crashing the polling loop. + """ + try: + import json as _json + + from app.models import IntegrationDirection, IntegrationType, UserIntegration + from app.utils.encryption import decrypt_value + + db = _get_db_session() + try: + integrations = ( + db.query(UserIntegration) + .filter( + UserIntegration.integration_type == IntegrationType.IMAP, + UserIntegration.direction == IntegrationDirection.SOURCE, + UserIntegration.is_active.is_(True), + ) + .all() + ) + logger.info("Processing %d IMAP UserIntegration(s)", len(integrations)) + for integ in integrations: + account_identifier = f"integration_{integ.owner_id}_{integ.id}" + try: + cfg = _json.loads(integ.config) if integ.config else {} + raw_creds = decrypt_value(integ.credentials) if integ.credentials else None + creds = _json.loads(raw_creds) if raw_creds else {} + + host = cfg.get("host") + port = int(cfg.get("port", 993)) + username = cfg.get("username") + password = creds.get("password") + use_ssl = cfg.get("use_ssl", True) + delete_after = cfg.get("delete_after_process", False) + gmail_labels = cfg.get("gmail_apply_labels", True) + + if not (host and username and password): + logger.warning( + "IMAP integration %d (owner %s) has incomplete config — skipping.", + integ.id, + integ.owner_id, + ) + continue + + pull_inbox( + mailbox_key=account_identifier, + host=host, + port=port, + username=username, + password=password, + use_ssl=use_ssl, + delete_after_process=delete_after, + owner_id=integ.owner_id, + gmail_apply_labels=gmail_labels, + ) + integ.last_used_at = datetime.now(timezone.utc) + integ.last_error = None + db.commit() + except Exception as exc: # noqa: BLE001 + error_msg = str(exc)[:_MAX_ERROR_LENGTH] + logger.error( + "Error pulling IMAP integration %d (owner %s): %s", + integ.id, + integ.owner_id, + error_msg, + ) + try: + integ.last_used_at = datetime.now(timezone.utc) + integ.last_error = error_msg + db.commit() + except Exception: # noqa: BLE001 + db.rollback() + finally: + db.close() + except Exception as exc: # noqa: BLE001 + logger.error("Failed to process IMAP UserIntegrations: %s", exc) + + +def check_and_pull_mailbox( + mailbox_key: str, + host: str | None, + port: int | None, + username: str | None, + password: str | None, + use_ssl: bool, + delete_after_process: bool, +): + """Validates config and invokes pulling from the mailbox if valid.""" + if not (host and port and username and password): + logger.warning(f"Mailbox {mailbox_key} is missing config, skipping.") + return + + logger.info(f"Checking mailbox: {mailbox_key}") + pull_inbox( + mailbox_key=mailbox_key, + host=host, + port=port, + username=username, + password=password, + use_ssl=use_ssl, + delete_after_process=delete_after_process, + ) + + +def pull_inbox( + mailbox_key, + host, + port, + username, + password, + use_ssl, + delete_after_process, + owner_id=None, + gmail_apply_labels=True, +): + """ + Connects to the IMAP inbox, fetches new unread emails from the last 3 days, + and processes attachments while preserving the original unread status. + + For Gmail: + - Attempts to select the localized All Mail folder. + - Runs an X-GM-RAW query: "in:anywhere in:unread newer_than:3d has:attachment". + + For non-Gmail mailboxes, it falls back to selecting the INBOX with a SINCE/UNSEEN filter. + + Args: + owner_id: Optional user identifier. When provided, ingested documents are + attributed to this user via ``process_document`` / ``convert_to_pdf``. + gmail_apply_labels: Whether to apply Gmail-specific labels and stars to + processed emails. Only relevant for Gmail hosts. Defaults to True. + """ + logger.info("Connecting to %s at %s:%s (SSL=%s)", mailbox_key, host, port, use_ssl) + processed_emails = load_processed_emails() + + try: + mail = imaplib.IMAP4_SSL(host, port) if use_ssl else imaplib.IMAP4(host, port) + mail.login(username, password) + + is_gmail_host = "gmail" in host.lower() + if is_gmail_host: + # For Gmail, try to select the localized All Mail folder. + all_mail_folder = find_all_mail_folder(mail) + if all_mail_folder: + logger.info("Using Gmail All Mail folder: %s", all_mail_folder) + mail.select(f'"{all_mail_folder}"') + else: + logger.warning("Gmail All Mail folder not found, falling back to INBOX.") + mail.select("INBOX") + # Use the X-GM-RAW query for Gmail. + raw_query = "in:anywhere in:unread newer_than:3d has:attachment" + status, search_data = mail.search(None, "X-GM-RAW", f'"{raw_query}"') + else: + # For non-Gmail, select INBOX and use SINCE/UNSEEN query. + mail.select("INBOX") + since_date = (datetime.now(timezone.utc) - timedelta(days=3)).strftime("%d-%b-%Y") + status, search_data = mail.search(None, f"(SINCE {since_date} UNSEEN)") + + if status != "OK": + logger.warning("Search failed on mailbox %s. Status=%s", mailbox_key, status) + mail.close() + mail.logout() + return + + msg_numbers = search_data[0].split() + logger.info("Found %d unread emails in %s.", len(msg_numbers), mailbox_key) + + for num in msg_numbers: + status, msg_data = mail.fetch(num, "(RFC822)") + if status != "OK": + logger.warning("Failed to fetch message %s in %s. Status=%s", num, mailbox_key, status) + continue + + raw_email = msg_data[0][1] + email_message = email.message_from_bytes(raw_email) + msg_id = email_message.get("Message-ID") + + if not msg_id: + logger.warning("Skipping email without Message-ID in %s", mailbox_key) + continue + + if msg_id in processed_emails: + logger.info("Skipping already processed email %s in %s", msg_id, mailbox_key) + continue + + # For Gmail, check if the email already has the "Ingested" label. + if is_gmail_host and gmail_apply_labels: + if email_already_has_label(mail, num, "Ingested"): + logger.info("Skipping email %s in %s, already labeled 'Ingested'.", msg_id, mailbox_key) + continue + + # Process attachments (and convert non-PDF files). + # We call the function without assigning its return value since it is not used. + fetch_attachments_and_enqueue(email_message, owner_id=owner_id) + + if settings.imap_readonly_mode: + logger.info("Readonly mode: skipping mailbox modifications for %s in %s", msg_id, mailbox_key) + else: + if is_gmail_host and gmail_apply_labels: + mark_as_processed_with_star(mail, num) + mark_as_processed_with_label(mail, num, label="Ingested") + + if delete_after_process: + logger.info("Deleting message %s from %s", num.decode(), mailbox_key) + mail.store(num, "+FLAGS", "\\Deleted") + else: + mail.store(num, "-FLAGS", "\\Seen") + + processed_emails[msg_id] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S") + save_processed_emails(processed_emails) + + if not settings.imap_readonly_mode and delete_after_process: + mail.expunge() + + mail.close() + mail.logout() + logger.info("Finished processing mailbox %s", mailbox_key) + + except Exception as e: + logger.exception("Error pulling mailbox %s: %s", mailbox_key, e) + + +def fetch_attachments_and_enqueue(email_message, owner_id: str | None = None): + """ + Extracts attachments from the email and processes only allowed file types. + + Files are accepted if either: + 1. They have a MIME type from the ALLOWED_MIME_TYPES set, OR + 2. They have a '.pdf' file extension (regardless of MIME type) + + Allowed file types include: + - PDF: application/pdf or *.pdf extension + - Microsoft Office files: + - Word: application/msword, + application/vnd.openxmlformats-officedocument.wordprocessingml.document + - Excel: application/vnd.ms-excel, + application/vnd.openxmlformats-officedocument.spreadsheetml.sheet + - PowerPoint: application/vnd.ms-powerpoint, + application/vnd.openxmlformats-officedocument.presentationml.presentation + - Other meaningful attachments: + - Plain text: text/plain + - CSV: text/csv + - Rich Text Format: application/rtf, text/rtf + + If the attachment is a PDF (by extension or MIME type), it is enqueued for upload; + any other allowed file is enqueued for conversion to PDF. + + Args: + email_message: The parsed email message to extract attachments from. + owner_id: Optional user identifier forwarded to ``process_document`` / + ``convert_to_pdf`` for multi-tenant attribution. + + Returns True if at least one allowed attachment was processed. + """ + has_attachment = False + for part in email_message.walk(): + if part.get_content_maintype() == "multipart": + continue + + filename = part.get_filename() + if not filename: + continue + + # Check if it's a PDF file by extension, regardless of MIME type + is_pdf_by_extension = filename.lower().endswith(".pdf") + + mime_type = part.get_content_type() + file_ext = os.path.splitext(filename)[1].lower() + # Accept file if it has an allowed MIME type, an allowed extension, OR is a PDF by extension + if mime_type not in ALLOWED_MIME_TYPES and file_ext not in ALLOWED_EXTENSIONS and not is_pdf_by_extension: + logger.info("Skipping attachment %s with MIME type %s", filename, mime_type) + continue + + file_path = os.path.join(settings.workdir, filename) + with open(file_path, "wb") as f: + f.write(part.get_payload(decode=True)) + + # If it's a PDF by MIME type or extension, process it directly + if mime_type == "application/pdf" or is_pdf_by_extension: + process_document.delay(file_path, owner_id=owner_id) + logger.info("Enqueued PDF for upload: %s (MIME: %s)", filename, mime_type) + elif mime_type in ALLOWED_MIME_TYPES: + # Other allowed files are sent for conversion + convert_to_pdf.delay(file_path, owner_id=owner_id) + logger.info("Enqueued file for conversion to PDF: %s", filename) + + has_attachment = True + return has_attachment + + +def email_already_has_label(mail, msg_id, label="Ingested"): + """ + Checks if the given message (msg_id) has the specified Gmail label. + Returns True if the label is found, False otherwise. + """ + try: + # Convert msg_id to bytes if it's an integer + if isinstance(msg_id, int): + msg_id = str(msg_id).encode() + + label_status, label_data = mail.fetch(msg_id, "(X-GM-LABELS)") + if label_status == "OK" and label_data and len(label_data) > 0: + raw_labels = label_data[0][1].decode("utf-8", errors="ignore") + if label in raw_labels: + return True + except Exception as e: + logger.error("Failed to fetch labels for msg_id=%s: %s", msg_id, e) + return False + + +def mark_as_processed_with_star(mail, msg_id): + """Stars the email in Gmail.""" + try: + mail.store(msg_id, "+FLAGS", "\\Flagged") + logger.info("Email %s starred in Gmail.", msg_id) + except Exception as e: + logger.error("Failed to star email %s: %s", msg_id, e) + + +def mark_as_processed_with_label(mail, msg_id, label="Ingested"): + """Adds a custom label to the email in Gmail.""" + try: + mail.store(msg_id, "+X-GM-LABELS", label) + logger.info("Email %s labeled '%s' in Gmail.", msg_id, label) + except Exception as e: + logger.error("Failed to label email %s with %s: %s", msg_id, label, e) + + +def find_all_mail_folder(mail): + """ + Attempts to select the Gmail All Mail folder using known localized names. + Falls back to using XLIST if needed. + Returns the folder name if found, otherwise None. + """ + COMMON_ALL_MAIL_NAMES = [ + "[Gmail]/Alle Nachrichten", + "[Gmail]/All Mail", + "[Gmail]/Todos", + "[Gmail]/Tutte le mail", + "[Gmail]/Tous les messages", + ] + for candidate in COMMON_ALL_MAIL_NAMES: + status, _ = mail.select(f'"{candidate}"', readonly=True) + if status == "OK": + return candidate + + capabilities = get_capabilities(mail) + if "XLIST" in capabilities: + candidate = find_all_mail_xlist(mail) + if candidate: + return candidate + return None + + +def get_capabilities(mail): + """Returns a list of capabilities supported by the IMAP server.""" + typ, data = mail.capability() + if typ == "OK" and data: + caps = data[0].decode("utf-8", errors="ignore").upper().split() + return caps + return [] + + +def find_all_mail_xlist(mail): + """ + Uses XLIST to discover the mailbox flagged as All Mail. + Returns the folder name if found, otherwise None. + """ + tag = mail._new_tag().decode("ascii") + command_str = f'{tag} XLIST "" "*"' + mail.send((command_str + "\r\n").encode("utf-8")) + + all_mail_folder = None + while True: + line = mail.readline() + if not line: + break + line_str = line.decode("utf-8", errors="ignore").strip() + if line_str.upper().startswith("* XLIST ") and "\\ALLMAIL" in line_str.upper(): + match = re.search(r'"([^"]+)"$', line_str) + if match: + candidate = match.group(1) + logger.info("Found All Mail folder via XLIST: %s", candidate) + all_mail_folder = candidate + if line_str.startswith(tag): + break + return all_mail_folder