gyt/opendan
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #48 from alexsunxl/MVP

Browse Source 
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
parent 7e1f7173b4 c6e18f96a9
commit 213e7b5b9b
1 changed files with 100 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/service/spider/email_spider.py
+65 -52
View File
@@ -5,91 +5,106 @@ import logging
import mailparser import mailparser
import hashlib import hashlib
import json import json
import base64
class EmailSpider:
def __init__(self):
# logger config # logger config
logger = logging.getLogger('email spider') self.logger = logging.getLogger('email spider')
logger.setLevel(logging.DEBUG) self.logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler() ch = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s [%(name)s] [%(levelname)s] %(message)s') formatter = logging.Formatter('%(asctime)s [%(name)s] [%(levelname)s] %(message)s')
ch.setFormatter(formatter) ch.setFormatter(formatter)
logger.addHandler(ch) 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')
)
client.login(self.config.get('EMAIL_ADDRESS'), self.config.get('EMAIL_PASSWORD'))
return client 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 # get email uid list
email_list = data[0].split() email_list = data[0].split()
logger.info(f"got {len(email_list)} emails") self.logger.info(f"got {len(email_list)} emails")
for uid in email_list: for uid in email_list:
logger.info(f"read email uid {uid}") if self.check_email_saved(uid):
if check_email_saved(client, uid): self.logger.info(f"email uid {uid} already saved")
logger.info(f"email uid {uid} already saved")
continue
else: else:
read_and_save_email(client, uid) self.read_and_save_email(uid)
logger.info(f"email uid {uid} saved") self.logger.info(f"email uid {uid} saved")
def read_and_save_email(self, uid: str):
def read_and_save_email(client: imaplib.IMAP4_SSL, 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):
# check only need to check email header, not need to download email body
def check_email_saved(client: imaplib.IMAP4_SSL, uid: str):
message_parts = "(BODY[HEADER])" message_parts = "(BODY[HEADER])"
_, 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"check email subject [{mail.subject}]") self.logger.info(f"[{uid}]check email subject [{mail.subject}]")
dir = get_local_dir_name(mail) dir = self.get_local_dir_name(mail)
logger.info(f"check email saved {dir}") self.logger.info(f"check email saved {dir}")
file = f"{dir}/email.txt" file = f"{dir}/email.txt"
if os.path.exists(file): if os.path.exists(file):
return True return False
return False return False
# save email attachment(images)
def save_email_attachment(self, mail: mailparser.MailParser, email_dir: str):
for attachment in mail.attachments:
if attachment['mail_content_type'] in ['image/png', 'image/jpeg', 'image/gif']:
print('current mail have image attachment')
img_dir = f"{email_dir}/image"
if not os.path.exists(img_dir):
os.makedirs(img_dir)
filename = attachment['filename']
filefullname = f"{img_dir}/{filename}"
image_data = attachment['payload']
try:
image_data = base64.b64decode(image_data)
except base64.binascii.Error:
image_data = image_data.encode()
with open(filefullname, 'wb') as f:
f.write(image_data)
self.logger.info(f"save email image {filename} success")
# save email to local file by each folder def save_email(self, mail: mailparser.MailParser):
def save_email(mail: mailparser.MailParser): dir = f"{self.config.get('LOCAL_DIR')}/{self.config.get('EMAIL_ADDRESS')}"
# create email account dir
dir = f"{config.get('LOCAL_DIR')}/{config.get('EMAIL_ADDRESS')}"
if not os.path.exists(dir): if not os.path.exists(dir):
os.makedirs(dir) os.makedirs(dir)
# create email local dir email_dir = self.get_local_dir_name(mail)
email_dir = get_local_dir_name(mail) self.logger.info(f"save email to {email_dir}")
logger.info(f"save email to {email_dir}")
if not os.path.exists(email_dir): if not os.path.exists(email_dir):
os.makedirs(email_dir) os.makedirs(email_dir)
# save email content and meta info
with open(f"{email_dir}/email.txt", "w") as f: with open(f"{email_dir}/email.txt", "w") as f:
f.write(mail.body) f.write(mail.body)
with open(f"{email_dir}/meta.json", "w", encoding='utf-8') as f: with open(f"{email_dir}/meta.json", "w", encoding='utf-8') as f:
@@ -97,13 +112,11 @@ def save_email(mail: mailparser.MailParser):
if 'body' in mail_dict: if 'body' in mail_dict:
del mail_dict['body'] del mail_dict['body']
json.dump(mail_dict, f, ensure_ascii=False, indent=4) json.dump(mail_dict, f, ensure_ascii=False, indent=4)
logger.info(f"save email meta info {f.name}") 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)