gitea-admin/gh-christianlouis-docuelevate
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update imap_tasks.py

Browse Source 
Refactor IMAP Tasks: Filter Allowed Attachment Types & Improve Code Quality

- Filter attachments to process only office files, PDFs, plain text, CSV, and RTF,
  explicitly skipping common image MIME types.
- Use an explicit elif for clarity in attachment processing logic.
- Remove dead code by not reassigning the result of fetch_attachments_and_enqueue.
- Reformat log messages, docstrings, and break long lines to comply with PEP 8.
- Add clarifying comments for excluded image MIME types.
This commit is contained in:
Christian Krakau-Louis
2025-03-13 02:08:21 +01:00
committed by GitHub
parent 473ae3a20a
commit 4c95e6ec6c
1 changed files with 45 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/tasks/imap_tasks.py
+45 -33
View File
@@ -141,7 +141,8 @@ def check_and_pull_mailbox(
)
def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_after_process):
def pull_inbox(mailbox_key, host, port, username, password, use_ssl,
delete_after_process):
"""
Connects to the IMAP inbox, fetches new unread emails from the last 3 days,
and processes attachments while preserving the original unread status.
@@ -152,7 +153,8 @@ def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_afte
For non-Gmail mailboxes, it falls back to selecting the INBOX with a SINCE/UNSEEN filter.
"""
logger.info(f"Connecting to {mailbox_key} at {host}:{port} (SSL={use_ssl})")
logger.info("Connecting to %s at %s:%s (SSL=%s)",
mailbox_key, host, port, use_ssl)
processed_emails = load_processed_emails()
try:
@@ -164,7 +166,7 @@ def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_afte
# For Gmail, try to select the localized All Mail folder.
all_mail_folder = find_all_mail_folder(mail)
if all_mail_folder:
logger.info(f"Using Gmail All Mail folder: {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.")
@@ -175,22 +177,25 @@ def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_afte
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")
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(f"Search failed on mailbox {mailbox_key}. Status={status}")
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(f"Found {len(msg_numbers)} unread emails in {mailbox_key}.")
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(f"Failed to fetch message {num} in {mailbox_key}. Status={status}")
logger.warning("Failed to fetch message %s in %s. Status=%s",
num, mailbox_key, status)
continue
raw_email = msg_data[0][1]
@@ -198,21 +203,23 @@ def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_afte
msg_id = email_message.get("Message-ID")
if not msg_id:
logger.warning(f"Skipping email without Message-ID in {mailbox_key}")
logger.warning("Skipping email without Message-ID in %s", mailbox_key)
continue
if msg_id in processed_emails:
logger.info(f"Skipping already processed email {msg_id} in {mailbox_key}")
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:
if email_already_has_label(mail, num, "Ingested"):
logger.info(f"Skipping email {msg_id} in {mailbox_key}, already labeled 'Ingested'.")
logger.info("Skipping email %s in %s, already labeled 'Ingested'.",
msg_id, mailbox_key)
continue
# Process attachments (and convert non-PDF files)
has_attachment = fetch_attachments_and_enqueue(email_message)
# 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)
if is_gmail_host:
mark_as_processed_with_star(mail, num)
@@ -222,7 +229,7 @@ def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_afte
save_processed_emails(processed_emails)
if delete_after_process:
logger.info(f"Deleting message {num.decode()} from {mailbox_key}")
logger.info("Deleting message %s from %s", num.decode(), mailbox_key)
mail.store(num, "+FLAGS", "\\Deleted")
else:
mail.store(num, "-FLAGS", "\\Seen")
@@ -232,10 +239,10 @@ def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_afte
mail.close()
mail.logout()
logger.info(f"Finished processing mailbox {mailbox_key}")
logger.info("Finished processing mailbox %s", mailbox_key)
except Exception as e:
logger.exception(f"Error pulling mailbox {mailbox_key}: {e}")
logger.exception("Error pulling mailbox %s: %s", mailbox_key, e)
def fetch_attachments_and_enqueue(email_message):
@@ -245,20 +252,26 @@ def fetch_attachments_and_enqueue(email_message):
Allowed file types include:
- PDF: application/pdf
- 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
- 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, it is enqueued for upload.
- For any other allowed file, a conversion task is enqueued to convert it to PDF.
Attachments not in this list are skipped. Common image MIME types such as
image/jpeg, image/png, image/gif, image/bmp, image/tiff, and image/webp are
intentionally excluded.
If the attachment is a PDF, it is enqueued for upload; any other allowed file
is enqueued for conversion to PDF.
Returns True if at least one allowed attachment was processed.
"""
# Define allowed MIME types for office files and other meaningful attachments.
ALLOWED_MIME_TYPES = {
"application/pdf",
"application/msword",
@@ -282,10 +295,10 @@ def fetch_attachments_and_enqueue(email_message):
if not filename:
continue
# Check the MIME type of the attachment.
mime_type = part.get_content_type()
if mime_type not in ALLOWED_MIME_TYPES:
logger.info(f"Skipping attachment {filename} with MIME type {mime_type}")
logger.info("Skipping attachment %s with MIME type %s",
filename, mime_type)
continue
file_path = os.path.join(settings.workdir, filename)
@@ -294,17 +307,16 @@ def fetch_attachments_and_enqueue(email_message):
if mime_type == "application/pdf":
upload_to_s3.delay(file_path)
logger.info(f"Enqueued PDF for upload: {filename}")
else:
logger.info("Enqueued PDF for upload: %s", filename)
elif mime_type in ALLOWED_MIME_TYPES:
# Enqueue conversion to PDF using the Gotenberg service.
convert_to_pdf.delay(file_path)
logger.info(f"Enqueued file for conversion to PDF: {filename}")
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.
@@ -317,7 +329,7 @@ def email_already_has_label(mail, msg_id, label="Ingested"):
if label in raw_labels:
return True
except Exception as e:
logger.error(f"Failed to fetch labels for msg_id={msg_id}: {e}")
logger.error("Failed to fetch labels for msg_id=%s: %s", msg_id, e)
return False
@@ -325,18 +337,18 @@ def mark_as_processed_with_star(mail, msg_id):
"""Stars the email in Gmail."""
try:
mail.store(msg_id, "+FLAGS", "\\Flagged")
logger.info(f"Email {msg_id} starred in Gmail.")
logger.info("Email %s starred in Gmail.", msg_id)
except Exception as e:
logger.error(f"Failed to star email {msg_id}: {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(f"Email {msg_id} labeled '{label}' in Gmail.")
logger.info("Email %s labeled '%s' in Gmail.", msg_id, label)
except Exception as e:
logger.error(f"Failed to label email {msg_id} with {label}: {e}")
logger.error("Failed to label email %s with %s: %s", msg_id, label, e)
def find_all_mail_folder(mail):
@@ -393,7 +405,7 @@ def find_all_mail_xlist(mail):
match = re.search(r'"([^"]+)"$', line_str)
if match:
candidate = match.group(1)
logger.info(f"Found All Mail folder via XLIST: {candidate}")
logger.info("Found All Mail folder via XLIST: %s", candidate)
all_mail_folder = candidate
if line_str.startswith(tag):
break