133 lines
5.1 KiB
Python
133 lines
5.1 KiB
Python
import concurrent.futures
|
|
import email
|
|
import hashlib
|
|
import imaplib
|
|
import time
|
|
from email.header import decode_header
|
|
from email.utils import parsedate_to_datetime
|
|
from json import JSONEncoder
|
|
from pathlib import Path
|
|
|
|
import chardet
|
|
|
|
from iarchiver.email import extract_emails
|
|
|
|
|
|
def md5_chunked(data: bytes, size: int = 1024):
|
|
m = hashlib.md5()
|
|
for i in range(0, len(data), size):
|
|
m.update(data[i:i + size])
|
|
return m.hexdigest()
|
|
|
|
|
|
class FileAttachment:
|
|
def __init__(self, file_name: str, file_hash: str):
|
|
self.filename = file_name
|
|
self.hash = file_hash
|
|
|
|
def to_dict(self):
|
|
return {'filename': self.filename, 'hash': self.hash}
|
|
|
|
|
|
class FileAttachmentEncoder(JSONEncoder):
|
|
def default(self, o):
|
|
if isinstance(o, FileAttachment):
|
|
return o.to_dict()
|
|
return super().default(o)
|
|
|
|
|
|
class MailConnection:
|
|
def __init__(self, host: str, username: str, password: str, attachments_dir: Path):
|
|
self.mail = imaplib.IMAP4_SSL(host)
|
|
self.mail.login(username, password)
|
|
self.attachments_dir = attachments_dir.expanduser().absolute().resolve()
|
|
self.folder_structure = {}
|
|
|
|
def load_folders(self):
|
|
folders = [tuple(f.decode().split(' "/" ')[1].replace('"', '').split('/')) for f in self.mail.list()[1]]
|
|
folder_structure = {}
|
|
for f in folders:
|
|
if not folder_structure.get(f[0]):
|
|
folder_structure[f[0]] = []
|
|
if len(f) > 1:
|
|
folder_structure[f[0]].append(f[1])
|
|
self.folder_structure = folder_structure
|
|
return self.folder_structure
|
|
|
|
def __fetch_email(self, i):
|
|
result, data = self.mail.uid('fetch', str(i), '(BODY[])') # fetch the raw email
|
|
if data[0] is None:
|
|
return
|
|
raw_email_bytes = data[0][1]
|
|
detected = chardet.detect(raw_email_bytes)
|
|
encoding = detected['encoding']
|
|
if not encoding:
|
|
encoding = 'utf-8'
|
|
raw_email = raw_email_bytes.decode(encoding, errors='replace')
|
|
email_message = email.message_from_string(raw_email)
|
|
date_header = email_message['Date']
|
|
date = parsedate_to_datetime(date_header)
|
|
unix_timestamp = int(time.mktime(date.timetuple()))
|
|
|
|
from_header = ', '.join(extract_emails(email_message['From']))
|
|
to_header = ', '.join(extract_emails(email_message['To']))
|
|
if '@' not in to_header:
|
|
to_header = email_message['To']
|
|
|
|
subject_header = email_message['Subject']
|
|
if subject_header:
|
|
subject = decode_header(subject_header)[0][0]
|
|
if isinstance(subject, bytes):
|
|
try:
|
|
detected = chardet.detect(subject)
|
|
encoding = detected['encoding']
|
|
if not encoding:
|
|
encoding = 'utf-8'
|
|
subject = subject.decode(encoding, errors='replace')
|
|
except UnicodeDecodeError:
|
|
subject = subject.decode('utf-8')
|
|
else:
|
|
return
|
|
|
|
attachments = []
|
|
if email_message.is_multipart():
|
|
for part in email_message.walk():
|
|
content_type = part.get_content_type()
|
|
content_disposition = str(part.get("Content-Disposition"))
|
|
if "attachment" in content_disposition:
|
|
filename = part.get_filename()
|
|
if filename:
|
|
# The filename of the file is the hash of its content, which should de-duplicate files.
|
|
filecontents = part.get_payload(decode=True)
|
|
filehash = md5_chunked(filecontents)
|
|
part.set_payload(filehash) # replace the attachment with its hash
|
|
filepath = self.attachments_dir / filehash
|
|
file_obj = FileAttachment(filename, filehash)
|
|
if not filepath.is_file():
|
|
with open(filepath, 'wb') as f:
|
|
f.write(filecontents)
|
|
attachments.append(file_obj)
|
|
raw_email_clean = email_message.as_string()
|
|
return unix_timestamp, to_header, from_header, subject, raw_email_clean, attachments
|
|
|
|
def fetch_folder(self, folder: str, search_criterion: str = 'ALL', max_threads: int = 1):
|
|
"""
|
|
Don't use multiple threads because most mail servers don't allow the client to multiplex.
|
|
"""
|
|
self.mail.select(f'"{folder}"')
|
|
result, data = self.mail.uid('search', None, search_criterion)
|
|
mail_ids = data[0]
|
|
id_list = mail_ids.split()
|
|
if not len(id_list):
|
|
# Empty folder
|
|
return
|
|
first_email_id = int(id_list[0])
|
|
latest_email_id = int(id_list[-1])
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor:
|
|
futures = {executor.submit(self.__fetch_email, i) for i in range(latest_email_id, first_email_id, -1)}
|
|
for future in concurrent.futures.as_completed(futures):
|
|
result = future.result()
|
|
if result is not None:
|
|
yield result
|