203 lines
7.5 KiB
Python
203 lines
7.5 KiB
Python
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
|
|
import re
|
|
|
|
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
|
|
|
|
# Helper function to sanitize filenames
|
|
def make_filename_safe(filename):
|
|
"""
|
|
Sanitizes a filename by:
|
|
- Removing any "file://" prefix.
|
|
- Replacing characters not allowed on most systems (< > : " / \ | ? *) with an underscore.
|
|
- Stripping leading/trailing whitespace.
|
|
|
|
Args:
|
|
filename (str): The original filename.
|
|
|
|
Returns:
|
|
str: A safe filename.
|
|
"""
|
|
# Remove any URL scheme if present
|
|
filename = filename.replace("file://", "")
|
|
# Replace illegal characters with underscore
|
|
safe_filename = re.sub(r'[<>:"/\\|?*]', '_', filename)
|
|
return safe_filename.strip()
|
|
|
|
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 (if you have a normalize_filename function, otherwise, use the safe version)
|
|
# Here we simply sanitize the filename.
|
|
filename = make_filename_safe(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:
|
|
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() |
