From 5a834ff690126d5c5ee109881fef7bfdbad6f10b Mon Sep 17 00:00:00 2001 From: Christian Krakau-Louis Date: Wed, 12 Feb 2025 01:54:41 +0100 Subject: [PATCH] added IMAP pull --- app/celery_worker.py | 11 +++ app/config.py | 18 ++++ app/tasks/imap_tasks.py | 203 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 232 insertions(+) create mode 100644 app/tasks/imap_tasks.py diff --git a/app/celery_worker.py b/app/celery_worker.py index 103635f0..6c15a964 100644 --- a/app/celery_worker.py +++ b/app/celery_worker.py @@ -19,6 +19,7 @@ from app.tasks.embed_metadata_into_pdf import embed_metadata_into_pdf from app.tasks.upload_to_dropbox import upload_to_dropbox from app.tasks.upload_to_paperless import upload_to_paperless from app.tasks.upload_to_nextcloud import upload_to_nextcloud +from app.tasks.imap_tasks import pull_all_inboxes celery.conf.task_routes = { "app.tasks.*": {"queue": "default"}, @@ -27,3 +28,13 @@ celery.conf.task_routes = { @celery.task def test_task(): return "Celery is working!" + +# If you want Celery Beat to run the poll task every minute, add: +from celery.schedules import crontab + +celery.conf.beat_schedule = { + "poll-inboxes-every-minute": { + "task": "app.tasks.imap_tasks.pull_all_inboxes", + "schedule": crontab(minute="*/1"), # every 1 minute + }, +} \ No newline at end of file diff --git a/app/config.py b/app/config.py index f10f8017..81c19321 100644 --- a/app/config.py +++ b/app/config.py @@ -26,6 +26,24 @@ class Settings(BaseSettings): paperless_ngx_url: str paperless_ngx_api_token: str paperless_host: str + + # IMAP 1 + imap1_host: Optional[str] = None + imap1_port: Optional[int] = 993 + imap1_username: Optional[str] = None + imap1_password: Optional[str] = None + imap1_ssl: bool = True + imap1_poll_interval_minutes: int = 5 + imap1_delete_after_process: bool = False + + # IMAP 2 + imap2_host: Optional[str] = None + imap2_port: Optional[int] = 993 + imap2_username: Optional[str] = None + imap2_password: Optional[str] = None + imap2_ssl: bool = True + imap2_poll_interval_minutes: int = 10 + imap2_delete_after_process: bool = False class Config: env_file = ".env" diff --git a/app/tasks/imap_tasks.py b/app/tasks/imap_tasks.py new file mode 100644 index 00000000..8c1a1d1e --- /dev/null +++ b/app/tasks/imap_tasks.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 + +import os +import time +import email +import imaplib +import tempfile +import logging +from datetime import datetime, timedelta + +from celery import shared_task +from app.config import settings + +# Import your pipeline's first-step task: +from app.tasks.upload_to_s3 import upload_to_s3 + +logger = logging.getLogger(__name__) + +# We'll keep track of the last poll time in a dictionary (mailbox -> datetime). +# This is in-memory only, so if you have multiple workers or restarts, it won't persist. +LAST_POLL_TIMES = { + "imap1": None, + "imap2": None, +} + + +@shared_task +def pull_all_inboxes(): + """ + This task is run periodically (e.g., every minute). + It checks each IMAP config, sees if it's time to poll (based on poll interval), + and if so, fetches attachments. + """ + logger.info("pull_all_inboxes - starting") + + # Check mailbox #1 + 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, + poll_interval=settings.imap1_poll_interval_minutes, + delete_after_process=settings.imap1_delete_after_process, + ) + + # Check mailbox #2 + 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, + poll_interval=settings.imap2_poll_interval_minutes, + delete_after_process=settings.imap2_delete_after_process, + ) + + logger.info("pull_all_inboxes - done") + + +def check_and_pull_mailbox( + mailbox_key: str, + host: str | None, + port: int | None, + username: str | None, + password: str | None, + use_ssl: bool, + poll_interval: int, + delete_after_process: bool, +): + """Check if it's time to poll this mailbox. If so, do it.""" + # If not configured, skip + if not (host and port and username and password): + logger.warning(f"Mailbox {mailbox_key} is missing config, skipping.") + return + + # If we never polled before, we can do it immediately. + last_poll = LAST_POLL_TIMES.get(mailbox_key) + now = datetime.utcnow() + + if last_poll is None or now - last_poll >= timedelta(minutes=poll_interval): + logger.info(f"Time to poll 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 + ) + LAST_POLL_TIMES[mailbox_key] = now + else: + logger.debug(f"Skipping {mailbox_key}, not due yet.") + + +def pull_inbox( + mailbox_key: str, + host: str, + port: int, + username: str, + password: str, + use_ssl: bool, + delete_after_process: bool +): + """ + Connect to the IMAP inbox, fetch UNSEEN messages, look for attachments, + and enqueue them to the pipeline. Then either delete or mark read. + """ + logger.info(f"Connecting to {mailbox_key} at {host}:{port} (SSL={use_ssl})") + + try: + if use_ssl: + mail = imaplib.IMAP4_SSL(host, port) + else: + mail = imaplib.IMAP4(host, port) + + mail.login(username, password) + mail.select("INBOX") + + # fetch only unseen messages + status, search_data = mail.search(None, "(UNSEEN)") + if status != "OK": + logger.warning(f"Search failed on mailbox {mailbox_key}. Status={status}") + mail.close() + mail.logout() + return + + msg_numbers = search_data[0].split() + logger.info(f"Found {len(msg_numbers)} new messages in {mailbox_key}.") + + for num in msg_numbers: + # fetch the full RFC822 message + status, msg_data = mail.fetch(num, "(RFC822)") + if status != "OK": + logger.warning(f"Failed to fetch message {num} in {mailbox_key}. Status={status}") + continue + + raw_email = msg_data[0][1] + email_message = email.message_from_bytes(raw_email) + + # Extract attachments and enqueue + fetch_attachments_and_enqueue(email_message) + + # Mark read or delete + if delete_after_process: + logger.info(f"Deleting message {num.decode()} from {mailbox_key}") + mail.store(num, "+FLAGS", "\\Deleted") + else: + logger.info(f"Marking message {num.decode()} as seen in {mailbox_key}") + mail.store(num, "+FLAGS", "\\Seen") + + if delete_after_process: + mail.expunge() + + mail.close() + mail.logout() + logger.info(f"Finished processing mailbox {mailbox_key}") + + except Exception as e: + logger.exception(f"Error pulling mailbox {mailbox_key}: {e}") + + +def fetch_attachments_and_enqueue(email_message): + """ + Iterate over parts of the email, extract PDF attachments, + save them to settings.workdir/tmp, and then enqueue + the pipeline with upload_to_s3.delay(). + """ + for part in email_message.walk(): + content_type = part.get_content_type() + filename = part.get_filename() + + if part.get_content_maintype() == 'multipart': + continue # skip container parts + if not filename: + continue # skip if there's no filename + + # Example: only PDF attachments + if content_type == "application/pdf": + logger.info(f"Found PDF attachment: {filename}") + + # Create a path in tmp with the original filename + file_path = os.path.join(settings.workdir, filename) + + # If a file with that name exists, we could rename it, but let's just overwrite + # or create a unique name. For simplicity: + with open(file_path, "wb") as f: + f.write(part.get_payload(decode=True)) + + logger.info(f"Saved attachment to {file_path}") + + # Now enqueue the pipeline's first step: + # E.g.: upload_to_s3 -> process_with_textract -> refine_text_with_gpt, etc. + upload_to_s3.delay(file_path) + + # If you prefer to *only* do the PDF embedding steps, you'd queue that first, etc. + # Or if you want to do the entire pipeline that you do in /process/ endpoint, + # you can replicate that chain here. + + # End fetch_attachments_and_enqueue