Merge pull request #48 from alexsunxl/MVP

The email spider module has added attachment image processing
This commit is contained in:
Liu Zhicong
2023-09-11 14:50:33 -07:00
committed by GitHub
+99 -86
View File
@@ -5,105 +5,118 @@ import logging
import mailparser import mailparser
import hashlib import hashlib
import json import json
import base64
# logger config class EmailSpider:
logger = logging.getLogger('email spider') def __init__(self):
logger.setLevel(logging.DEBUG) # logger config
ch = logging.StreamHandler() self.logger = logging.getLogger('email spider')
formatter = logging.Formatter('%(asctime)s [%(name)s] [%(levelname)s] %(message)s') self.logger.setLevel(logging.DEBUG)
ch.setFormatter(formatter) ch = logging.StreamHandler()
logger.addHandler(ch) formatter = logging.Formatter('%(asctime)s [%(name)s] [%(levelname)s] %(message)s')
ch.setFormatter(formatter)
self.logger.addHandler(ch)
# read config from toml file # read config from toml file
# and read from config config.local.toml if exists (config.local.toml is ignored by git) # and read from config config.local.toml if exists (config.local.toml is ignored by git)
config = toml.load('./rootfs/email/config.toml') self.config = toml.load('./rootfs/email/config.toml')
if os.path.exists('./rootfs/email/config.local.toml'): if os.path.exists('./rootfs/email/config.local.toml'):
config = toml.load('./rootfs/email/config.local.toml') self.config = toml.load('./rootfs/email/config.local.toml')
self.client = self.email_client()
# create email client def email_client(self) -> imaplib.IMAP4_SSL:
def email_client() -> imaplib.IMAP4_SSL: self.logger.info(f"read email config from {self.config.get('EMAIL_IMAP_SERVER')}")
logger.info(f"read email config from {config.get('EMAIL_IMAP_SERVER')}") client = imaplib.IMAP4_SSL(
client = imaplib.IMAP4_SSL(config.get('EMAIL_IMAP_SERVER')) host=self.config.get('EMAIL_IMAP_SERVER'),
client.login(config.get('EMAIL_ADDRESS'), config.get('EMAIL_PASSWORD')) port=self.config.get('EMAIL_IMAP_PORT')
return client )
client.login(self.config.get('EMAIL_ADDRESS'), self.config.get('EMAIL_PASSWORD'))
return client
def list_box(mail: imaplib.IMAP4_SSL): def list_box(self):
_, mailbox_list = mail.list() _, mailbox_list = self.client.list()
for mailbox in mailbox_list: for mailbox in mailbox_list:
print(mailbox.decode()) print(mailbox.decode())
def read_emails(client: imaplib.IMAP4_SSL, folder: str = 'INBOX', imap_keyword: str = "UNSEEN"): def read_emails(self, folder: str = 'INBOX', imap_keyword: str = "UNSEEN"):
client.select(folder) self.client.select(folder)
_, data = client.uid('search', None, imap_keyword) _, data = self.client.uid('search', None, imap_keyword)
# get email uid list
email_list = data[0].split()
logger.info(f"got {len(email_list)} emails")
for uid in email_list:
logger.info(f"read email uid {uid}")
if check_email_saved(client, uid):
logger.info(f"email uid {uid} already saved")
continue
else:
read_and_save_email(client, uid)
logger.info(f"email uid {uid} saved")
# get email uid list
email_list = data[0].split()
self.logger.info(f"got {len(email_list)} emails")
for uid in email_list:
if self.check_email_saved(uid):
self.logger.info(f"email uid {uid} already saved")
else:
self.read_and_save_email(uid)
self.logger.info(f"email uid {uid} saved")
def read_and_save_email(client: imaplib.IMAP4_SSL, uid: str): def read_and_save_email(self, uid: str):
message_parts = "(BODY.PEEK[])" message_parts = "(BODY.PEEK[])"
_, email_data = client.uid('fetch', uid, message_parts) _, email_data = self.client.uid('fetch', uid, message_parts)
mail = mailparser.parse_from_bytes(email_data[0][1]) mail = mailparser.parse_from_bytes(email_data[0][1])
logger.info(f"got email subject [{mail.subject}]") self.logger.info(f"got email subject [{mail.subject}]")
save_email(mail) self.save_email(mail)
def get_local_dir_name(mail: mailparser.MailParser) -> str: def get_local_dir_name(self, mail: mailparser.MailParser) -> str:
dir = f"{config.get('LOCAL_DIR')}/{config.get('EMAIL_ADDRESS')}" dir = f"{self.config.get('LOCAL_DIR')}/{self.config.get('EMAIL_ADDRESS')}"
name = f"{mail.subject}__{mail.date}" name = f"{mail.subject}__{mail.date}"
name = hashlib.md5(name.encode('utf-8')).hexdigest() name = hashlib.md5(name.encode('utf-8')).hexdigest()
return f"{dir}/{name}" return f"{dir}/{name}"
def check_email_saved(self, uid: str):
message_parts = "(BODY[HEADER])"
_, email_data = self.client.uid('fetch', uid, message_parts)
mail = mailparser.parse_from_bytes(email_data[0][1])
self.logger.info(f"[{uid}]check email subject [{mail.subject}]")
dir = self.get_local_dir_name(mail)
self.logger.info(f"check email saved {dir}")
file = f"{dir}/email.txt"
if os.path.exists(file):
return False
return False
# check only need to check email header, not need to download email body # save email attachment(images)
def check_email_saved(client: imaplib.IMAP4_SSL, uid: str): def save_email_attachment(self, mail: mailparser.MailParser, email_dir: str):
message_parts = "(BODY[HEADER])" for attachment in mail.attachments:
_, email_data = client.uid('fetch', uid, message_parts) if attachment['mail_content_type'] in ['image/png', 'image/jpeg', 'image/gif']:
mail = mailparser.parse_from_bytes(email_data[0][1]) print('current mail have image attachment')
logger.info(f"check email subject [{mail.subject}]") img_dir = f"{email_dir}/image"
dir = get_local_dir_name(mail) if not os.path.exists(img_dir):
logger.info(f"check email saved {dir}") os.makedirs(img_dir)
file = f"{dir}/email.txt" filename = attachment['filename']
if os.path.exists(file): filefullname = f"{img_dir}/{filename}"
return True image_data = attachment['payload']
return False try:
image_data = base64.b64decode(image_data)
except base64.binascii.Error:
# save email to local file by each folder image_data = image_data.encode()
def save_email(mail: mailparser.MailParser): with open(filefullname, 'wb') as f:
# create email account dir f.write(image_data)
dir = f"{config.get('LOCAL_DIR')}/{config.get('EMAIL_ADDRESS')}" self.logger.info(f"save email image {filename} success")
if not os.path.exists(dir):
os.makedirs(dir)
# create email local dir
email_dir = get_local_dir_name(mail)
logger.info(f"save email to {email_dir}")
if not os.path.exists(email_dir):
os.makedirs(email_dir)
# save email content and meta info
with open(f"{email_dir}/email.txt", "w") as f:
f.write(mail.body)
with open(f"{email_dir}/meta.json", "w", encoding='utf-8') as f:
mail_dict = json.loads(mail.mail_json)
if 'body' in mail_dict:
del mail_dict['body']
json.dump(mail_dict, f, ensure_ascii=False, indent=4)
logger.info(f"save email meta info {f.name}")
def save_email(self, mail: mailparser.MailParser):
dir = f"{self.config.get('LOCAL_DIR')}/{self.config.get('EMAIL_ADDRESS')}"
if not os.path.exists(dir):
os.makedirs(dir)
email_dir = self.get_local_dir_name(mail)
self.logger.info(f"save email to {email_dir}")
if not os.path.exists(email_dir):
os.makedirs(email_dir)
with open(f"{email_dir}/email.txt", "w") as f:
f.write(mail.body)
with open(f"{email_dir}/meta.json", "w", encoding='utf-8') as f:
mail_dict = json.loads(mail.mail_json)
if 'body' in mail_dict:
del mail_dict['body']
json.dump(mail_dict, f, ensure_ascii=False, indent=4)
self.logger.info(f"save email meta info {f.name}")
self.save_email_attachment(mail, email_dir)
if __name__ == "__main__": if __name__ == "__main__":
mail = email_client() spider = EmailSpider()
folder = 'INBOX' folder = 'INBOX'
# imap_keyword = "UNSEEN"
imap_keyword = "ALL" imap_keyword = "ALL"
read_emails(mail, folder, imap_keyword) spider.read_emails(folder, imap_keyword)