MailMigration/attachments.py

from email import message_from_bytes
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText

import os
import hashlib
from email.utils import parsedate_to_datetime
from pathlib import Path
from datetime import datetime

from config import get_config
from stats import add_statistic

# Global variable: Define which attachment types and sizes should be extracted
EXTRACT_ATTACHMENTS = get_config("extract_attachments", "enabled", False)
ATTACHMENT_WHITELIST = get_config("extract_attachments", "attachment_whitelist", [".pdf", ".zip", ".docx", ".xlsx"]) # Allowed file types
MIN_ATTACHMENT_SIZE = get_config("extract_attachments", "min_attachment_size", 0 * 1024) # Min size for attachments to have to be downloaded
MAX_ATTACHMENT_SIZE = get_config("extract_attachments", "max_attachment_size", 100 * 1024 * 1024)  # Max size per attachment in Bytes

download_folder = os.path.expanduser("~")+"/Downloads/"
EXTERNAL_STORAGE_PATH = get_config("extract_attachments", "external_storage_path", download_folder)  # Path to store extracted files

def normalize_filename(filename):
    """
    Normalizes the filename by removing any trailing query string or extra characters
    from the file extension. For example:
      "document.pdf?=" becomes "document.pdf"
    """
    if '?' in filename:
        filename = filename.split('?')[0]
    return filename.strip()

def get_normalized_extension(filename):
    """
    Returns the normalized file extension (in lower-case) for the given filename.
    """
    filename = normalize_filename(filename)
    return os.path.splitext(filename)[1].lower()


def should_extract_attachment(part):
    """Decide whether an attachment should be extracted based on type and size."""
    filename = part.get_filename()
    if not filename:
        return False
    
    # Normalize filename and extension.
    filename = normalize_filename(filename)
    file_ext = get_normalized_extension(filename)
    file_size = len(part.get_payload(decode=True))  # File size in bytes

    add_statistic("attachment_types", file_ext)

    return file_ext in ATTACHMENT_WHITELIST and file_size < MAX_ATTACHMENT_SIZE and file_size > MIN_ATTACHMENT_SIZE

def save_attachment(part, email_date_str):
    """
    Save an extracted attachment to disk in a month-based folder (e.g., "2023_10")
    and return a file:// URL pointing to the stored file.
    
    Args:
        part: The email attachment part.
        email_date_str: The email's sent date as a string 
                        (e.g., "Mon, 09 Oct 2023 12:34:56 -0400" or '"10-Feb-2025 07:56:34 +0100"')
    
    Returns:
        A file:// URL to the stored attachment, or None if filename is missing.
    """
    filename = part.get_filename()
    if not filename:
        return None

    # Normalize the filename to remove trailing query strings etc.
    filename = normalize_filename(filename)

    # Strip extra quotes from the date string, if present.
    date_str = email_date_str.strip('"')
    
    # Parse the email date string into a datetime object.
    try:
        email_date = parsedate_to_datetime(date_str)
    except Exception as e:
        try:
            # Fallback: try using datetime.strptime with the appropriate format.
            email_date = datetime.strptime(date_str, "%d-%b-%Y %H:%M:%S %z")
        except Exception as e2:
            print(f"Date parsing failed for: {date_str}")
            print(e2)
            email_date = datetime.now()

    # Create a folder name based on the email's date (e.g., "2023_10")
    month_folder = email_date.strftime("%Y_%m")
    storage_dir = os.path.join(EXTERNAL_STORAGE_PATH, month_folder)
    os.makedirs(storage_dir, exist_ok=True)

    # Generate a unique filename using a hash prefix to avoid collisions.
    hash_prefix = hashlib.md5(filename.encode()).hexdigest()[:8]
    storage_filename = f"{hash_prefix}_{filename}"
    storage_path = os.path.join(storage_dir, storage_filename)

    # Write the attachment data to disk.
    file_data = part.get_payload(decode=True)
    with open(storage_path, "wb") as f:
        f.write(file_data)

    # Create a file:// URL pointing to the stored file using pathlib.
    file_url = Path(storage_path).as_uri()
    return file_url

def extract_and_replace_attachments(raw_msg, email_date_str):
    """
    Extracts attachments meeting the criteria and replaces them with a link.
    Reassembles the email so that the body appears only once.

    Args:
        raw_msg (bytes): Raw email message in bytes.
        email_date_str: Sent date of the message (used for organizing storage).

    Returns:
        bytes: The modified email message (as bytes) with attachments replaced by links.
    """
    # Parse the raw message into an email object
    email_msg = message_from_bytes(raw_msg)
    
    # If the message isn't multipart, return it unchanged
    if not email_msg.is_multipart():
        return raw_msg

    # Create a new email container (multipart/mixed)
    new_email = MIMEMultipart("mixed")
    # Copy key headers from the original message
    for header in ["Subject", "From", "To", "Date"]:
        if email_msg[header]:
            new_email[header] = email_msg[header]

    # Initialize lists to store parts
    body_parts = []         # For text parts (the email body)
    keep_parts = []         # For attachments that should not be extracted
    attachment_links = []   # For links to extracted attachments

    # Walk through all parts (leaf nodes) of the email
    for part in email_msg.walk():
        # Skip container parts
        if part.is_multipart():
            continue

        # Check if the part has a filename (i.e. it's an attachment)
        filename = part.get_filename()
        if filename:
            # This part is an attachment.
            if should_extract_attachment(part):
                link = save_attachment(part, email_date_str)
                if link:
                    attachment_links.append(link)
            else:
                # Keep attachments that don't meet extraction criteria.
                keep_parts.append(part)
        else:
            # If there is no filename, assume it's a body (text) part.
            body_parts.append(part)

    # Choose a single body part: Prefer "text/plain" if available.
    selected_body = None
    for part in body_parts:
        if part.get_content_type() == "text/plain":
            selected_body = part
            break
    if not selected_body and body_parts:
        selected_body = body_parts[0]

    # Assemble the new email.
    if selected_body:
        new_email.attach(selected_body)
    # Attach any attachments that should be kept intact.
    for part in keep_parts:
        new_email.attach(part)
    # If any attachments were extracted, add a summary text part with links.
    if attachment_links:
        links_text = "\n\n[Attachments extracted and stored separately:]\n" + "\n".join(attachment_links)
        new_email.attach(MIMEText(links_text, "plain"))
    
    return new_email.as_bytes()